# A Constrained Optimization Perspective of Unrolled Transformers

arXiv: https://arxiv.org/abs/2601.17257 · HTML: https://mlsanity.com/paper/2601.17257/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 42 of its 66 references that are in the map.
- First posted on arXiv 24 Jan 2026.

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Deep Learning with Differential Privacy](https://mlsanity.com/paper/1607.00133/) (2016)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [A Closer Look at Memorization in Deep Networks](https://mlsanity.com/paper/1706.05394/) (2017)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [Algorithm Unrolling: Interpretable, Efficient Deep Learning for Signal and Image Processing](https://mlsanity.com/paper/1912.10557/) (2019)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Differentially Private Fine-tuning of Language Models](https://mlsanity.com/paper/2110.06500/) (2021)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Diffusion models as plug-and-play priors](https://mlsanity.com/paper/2206.09012/) (2022)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers](https://mlsanity.com/paper/2212.10559/) (2022)
- [Plug-and-Play Methods Provably Converge with Properly Trained Denoisers](https://mlsanity.com/paper/1905.05406/) (2019)
- [Privacy- and Utility-Preserving Textual Analysis via Calibrated Multivariate Perturbations](https://mlsanity.com/paper/1910.08902/) (2019)
- [Uncovering mesa-optimization algorithms in Transformers](https://mlsanity.com/paper/2309.05858/) (2023)
- [Theoretical Linear Convergence of Unfolded ISTA and its Practical Weights and Thresholds](https://mlsanity.com/paper/1808.10038/) (2018)
- [White-Box Transformers via Sparse Rate Reduction](https://mlsanity.com/paper/2306.01129/) (2023)
- [Deep Unfolding: Model-Based Inspiration of Novel Deep Architectures](https://mlsanity.com/paper/1409.2574/) (2014)
- [Lagrangian Duality for Constrained Deep Learning](https://mlsanity.com/paper/2001.09394/) (2020)
- [Graph Neural Networks Inspired by Classical Iterative Algorithms](https://mlsanity.com/paper/2103.06064/) (2021)
- [Probably Approximately Correct Constrained Learning](https://mlsanity.com/paper/2006.05487/) (2020)
- [Constrained Learning with Non-Convex Losses](https://mlsanity.com/paper/2103.05134/) (2021)

## Abstract (arXiv)

We introduce a constrained optimization framework for training transformers that behave like optimization descent algorithms. Specifically, we enforce layerwise descent constraints on the objective function and replace standard empirical risk minimization (ERM) with a primal-dual training scheme. This approach yields models whose intermediate representations decrease the loss monotonically in expectation across layers. We apply our method to both unrolled transformer architectures and conventional pretrained transformers on tasks of video denoising and text classification. Across these settings, we observe constrained transformers achieve stronger robustness to perturbations and maintain higher out-of-distribution generalization, while preserving in-distribution performance.