Title: A Constrained Optimization Perspective of Unrolled Transformers
Primary area: optimization
Keywords: constrained learning, unrolled neural networks, transformers

Abstract: 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 that constrained transformers achieve stronger robustness to perturbations and maintain higher out-of-distribution generalization, while preserving competitive in-distribution performance.

TL;DR: Transformers trained with monotonic descent constraints are more robust to out-of-distribution perturbations.