A Constrained Optimization Perspective of Unrolled Transformers

TL;DR

The paper addresses transformer fragility under perturbations and distribution shifts by enforcing layerwise descent of the loss using a constrained, dual-structured learning problem. The authors develop a primal--dual training algorithm, proving a Constrained Learning Theorem that bounds the duality gap and an asymptotic convergence result, and they establish OOD generalization guarantees under distribution shifts. Empirically, constrained transformers show improved robustness with minimal or no sacrifice to in-distribution performance across video denoising and text classification with perturbed embeddings, and extend to large-scale models such as Llama in finetuning scenarios. These findings suggest that imposing monotone layerwise descent is a practical and scalable strategy to enhance transformer robustness in real-world, perturbed settings.

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

Figures (16)

• Figure 1: Layerwise descent improves OOD robustness. Left: Test loss at each layer ($\downarrow$ lower is better). Constrained RoBERTa exhibits monotonic descent, unlike the unconstrained baseline. Right: Out-of-distribution accuracy under increasing embedding perturbation levels $\gamma$ ($\uparrow$ higher is better). As $\gamma$ grows, the constrained model degrades more gracefully and retains higher accuracy. Setting: RoBERTa ($L=24$) trained on IMDb, training $\gamma=0.2$.
• Figure 2: Video denoising error vs. test perturbation $\gamma$ (RMSE $\downarrow$ , lower is better). Columns are datasets, rows are architectures. Solid lines are constrained models; dashed lines are unconstrained. Each plot shows RMSE over increasing test perturbation levels ($\gamma$). All models were trained with perturbation $\gamma_\text{train}=0.13$.
• Figure 3: Effects of descent constraints on the OOD performance of (left, $\downarrow$ lower is better) unconstrained DUST and (right, $\downarrow$ lower is better) constrained DUST trained on UCSD with training $\gamma=0$. Deeper constrained models show improved OOD performance, while increasing depth does not benefit unconstrained models.
• Figure 4: Text classification accuracy vs. test perturbation (Accuracy $\uparrow$, higher is better). Columns are datasets, rows are architectures. Solid lines denote constrained models; dashed lines denote unconstrained. Each plot shows accuracy over increasing test perturbation levels ($\gamma$). All models were trained with perturbation $\gamma_{\text{train}}=0.8$.
• Figure 5: Ablation analysis of constrained RoBERTa on IMDb. Left: layerwise validation loss ($\downarrow$ lower is better), right: test accuracy across test perturbation levels $\uparrow$ higher is better. Increasing the constraint step size $\alpha$ induces monotonic descent along the layers, while improving OOD robustness.

Theorems & Definitions (3)

• proof
• proof
• proof