Technical Summary: A Constrained Optimization Perspective of Unrolled Transformers

Problem Statement

While recent literature has interpreted Transformer architectures as iterative descent algorithms solving optimization problems (e.g., Yang et al., 2022; De Weerdt et al., 2023), standard training via Empirical Risk Minimization (ERM) often fails to enforce the monotonic descent behavior expected of such optimizers. When trained with ERM, the intermediate representations of unrolled Transformers do not necessarily decrease the loss function monotonically across layers. This inconsistency arises because standard training does not guarantee that learned parameters satisfy the conditions required for energy minimization, nor does it account for the bilevel nature of unrolled architecture training. Consequently, these models may exhibit non-monotonic loss patterns and reduced robustness to input perturbations and distribution shifts.

Methodology

The authors propose a framework that treats Transformer training as a constrained learning problem. Instead of minimizing the average loss over the entire network output, the method enforces that the expected loss decreases monotonically at every layer.

1. Constrained Formulation

The core contribution is the formulation of a constrained optimization problem where the output of each layer $l$ is required to reduce the expected loss $f$ by a factor $(1 - \alpha_l)$ relative to the previous layer $l-1$:
$$\begin{aligned} T^* = \arg\min_{T} \quad & \mathbb{E}[f(X, \Phi(X; T))] \\ \text{subject to} \quad & \mathbb{E}[f(X, \Phi_l(X; T))] \leq (1 - \alpha_l) \mathbb{E}[f(X, \Phi_{l-1}(X; T))], \quad \forall l \end{aligned}$$
Here, $\Phi_l$ represents the output of layer $l$, and $\alpha_l$ is a stepsize schedule. This formulation forces the Transformer to behave like a descent algorithm, ensuring intermediate representations progressively improve the objective.

2. Primal-Dual Training Algorithm

To solve this non-convex constrained problem, the authors employ a primal-dual training scheme in the dual domain. They construct a Lagrangian function:
$$\mathcal{L}(T, \lambda) = \mathbb{E}[f(X, \Phi(X; T))] + \sum_{l=1}^L \lambda_l \mathbb{E}[f(X, \Phi_l(X; T)) - (1 - \alpha_l)f(X, \Phi_{l-1}(X; T))]$$
The training procedure (Algorithm 1) alternates between:

• Primal Update: Minimizing the empirical Lagrangian with respect to Transformer parameters $T$ using stochastic gradient descent.
• Dual Update: Maximizing the dual function with respect to Lagrange multipliers $\lambda$ using projected gradient ascent.

The authors leverage recent results showing that deep neural networks exhibit small duality gaps (Chamon et al., 2023), allowing the dual problem to yield near-optimal and near-feasible solutions. To ensure strict feasibility, a resilient relaxation with a slack variable and quadratic penalty is introduced.

Key Contributions

The paper outlines four primary contributions:

1. Constrained Learning Formulation: The authors formulate a learning problem where Transformer layers are explicitly constrained to reduce the expected loss monotonically, diverging from standard ERM.
2. Dual Training Algorithm & Convergence Guarantees: They develop a primal-dual training algorithm and provide theoretical guarantees (Theorem 2) that the constrained unrolled Transformer converges asymptotically to the optimal value of the statistical loss. They further establish Out-of-Distribution (OOD) generalization guarantees (Theorem 4), proving that the descent properties are maintained under distribution shifts.
3. Robustness to Perturbed Embeddings: In text classification tasks (IMDb, MNLI), the authors demonstrate that constrained models maintain higher accuracy under Gaussian, uniform, and discrete perturbations to token embeddings compared to unconstrained baselines.
4. Robustness in Video Denoising: Applied to video denoising (CUHK Avenue, UCSD, ShanghaiTech), the method yields models that are more robust to OOD noise levels, measured by Root Mean Squared Error (RMSE), while preserving in-distribution performance.

Experimental Results

The authors evaluated their method on both conventional pre-trained Transformers (DistilBERT, RoBERTa, Llama 3.1) and unrolled architectures (Deep Unfolded Sequential Transformer - DUST, Unrolled Transformer - UT).

• Monotonic Descent: Constrained models exhibited smooth, monotonic decreases in layer-wise loss, whereas unconstrained baselines showed non-monotonic patterns with spikes (e.g., a spike at layer 19 in RoBERTa).
• OOD Robustness:
  • Text Classification: As perturbation levels ($\gamma$) increased, constrained models degraded more gracefully. For instance, on RoBERTa with IMDb, the constrained model retained significantly higher accuracy under high noise levels compared to the unconstrained version.
  • Video Denoising: Constrained models consistently achieved lower RMSE under high test perturbation levels ($\gamma \geq 1.0$) across various datasets and architectures.
  • LLM Fine-tuning: On Llama 3.1 8B, the constrained model maintained in-distribution performance (50% win rate on AlpacaEval with no noise) while significantly outperforming the unconstrained model under noise (69.93% win rate with $\gamma=1.5$).
• Trade-offs: The method generally preserved in-distribution performance with negligible degradation (often < 0.3%), suggesting that enforcing descent constraints does not sacrifice standard accuracy for robustness.

Significance and Claims

The paper claims that enforcing descent constraints aligns the internal dynamics of Transformers with the stability properties of classical optimization algorithms. This alignment provides a theoretical and empirical basis for improved robustness against input perturbations and distribution shifts.

The authors position this work as a bridge between algorithmic unrolling (where networks solve optimization problems) and constrained learning. They argue that by explicitly training Transformers to descend, one can achieve models that are inherently more robust to noise, a property particularly valuable for applications involving local differential privacy, where input perturbations are necessary for user privacy. The work suggests that the "unrolled" perspective is not merely a theoretical interpretation but a practical training objective that yields tangible improvements in generalization and robustness without requiring architectural changes.