{'fraction_non_alphanumeric': 0.04253745651757101, 'fraction_numerical': 0.0121887753704566, 'mean_word_length': 4.621247113163972, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Transformers from an Optimization Perspective

Deep learning models such as the Transformer are often constructed by heuristics and experience. To provide a complementary foundation, in this work we study the following problem: Is it possible to find an energy function underlying the Transformer model, such that descent steps along this energy correspond with the Transformer forward pass? By finding such a function, we can reinterpret Transformers as the unfolding of an interpretable optimization process across iterations. This unfolding perspective has been frequently adopted in the past to elucidate more straightforward deep models such as MLPs and CNNs; however, it has thus far remained elusive obtaining a similar equivalence for more complex models with self-attention mechanisms like the Transformer. To this end, we first outline several major obstacles before providing companion techniques to at least partially address them, demonstrating for the first time a close association between energy function minimization and deep layers with self-attention. This interpretation contributes to our intuition and understanding of Transformers, while potentially laying the ground-work for new model designs.

We take significant steps towards extending the unfolded optimization perspective to a full Transformer stack via the following contributions: • After providing background details and existing examples of unfolded optimization, Section 2 formalizes the four key challenges to mapping Transformer layers to optimization steps descending a well-specified energy function. These include: (i) handling self-attention, (ii) integrating heterogeneous Transformer layer modules (i.e., self-attention and FFN), (iii) accounting for non-linear activations, and (iv) allowing for arbitrary/unstructured weight matrices.
• Later, Section 3 derives an energy function that closely reproduces self-attention per Challenge (i), Section 4 rigorously derives convergence results which demonstrate how to approximately handle Challenge (ii) via a novel form of alternating minimization, and Section 5 provides explicit details for solving Challenge (iii) via proximal methods; consideration of Challenge (iv) is deferred to the supplementary along with complementary supporting analyses.
• The above analysis culminates in an energy function that is minimized by Transformer layers when certain technical conditions hold. To provide further support, Section 6 empirically demonstrates that indeed this energy is minimized with real-world data even when technical conditions are difficult to formally verify.
Collectively, these results provide a complementary foundation for understanding Transformer models, and suggest plausible avenues for future enhancements.

Brief Intro to Deep Architectures Formed from Unfolded Optimization
The basic idea of unfolded optimization perspective is to create a one-to-one correspondence between the layers of a deep model and the iterations of some algorithmic process designed to minimize a parameterized, model-specific energy function. In this way, the forward pass can be viewed as computing an approximation to the minimum energy solution. Meanwhile, assuming each forward iteration is differentiable w.r.t. the energy function parameters, we can pass gradients of some meta objective function (e.g., classification, regression, etc.) through the approximate minimization steps for the backward pass, which forms an interpretable bilevel-optimization process [10,13,23]. Overall then, unfolded optimization produces minimal energy representations, whereby the specific structure of the energy function is trainable per the relevant downstream application domain.

Mathematical Formulation
The bilevel optimization process described above can be expressed more formally as: Here X ∈ R n×d is the model input (e.g., input sentence to a language model), where n is the number of tokens and d is the token dimensionality. We refer to Y ∈ R n×d as the hidden representation of the model, where Y * is the (ideal) model output (e.g. word representation output by a Transformer encoder) that minimizes to lower-level energy function 2 E : R n×d → R whose form is dictated by parameters W (e.g., these could also correspond with the trainable parameters in a Transformer model). And finally, is a meta-loss function (e.g., the loss function used to train a Transformer) while ψ : R n×d → R n×d is an arbitrary transformation that converts hidden representations to the model output, which can also potentially have its own learnable parameters. Assuming ∂Y * (W, X)/∂W and ∂ /∂Y * (W, X) are computable, where the former involves passing gradients through the iterative steps used to approximate Y * (W, X), then the combined bilevel system can be optimized w.r.t. W .
Note that hereinafter if the energy function does not include inter-token interactions, we view it as a function of vectors, and we use lower-case characters y and x to represent an arbitrary row of Y and X for simplicity. Also, when clear from context, we omit writing the parameters of a function, for example we may write E(Y ) to reference E(Y ; W, X).

Related Work and Limitations
In this section we present related work on unfolded optimization, and in particular, analyze some of the relevant limitations that serve as motivation for our efforts and underscore the challenges involved.
Optimization Induced Feed-Forward Networks A wide variety of prior work has investigated feedforward structures from an unfolded optimization perspective [13,15,19,21,41,47]. Here we take a closer look at [47] which is based on proximal methods. The basic energy function in [47] is where 1 is a vector with all entries being 1,σ * is the convex conjugate [38] 3 ofσ : a → a 0 σ(r)dr, σ is an activation function and f (x) is some transformation of input prepresentation x. The proximal operator [11] of (3) is Note that although in (4) W can in principle be arbitrary matrix as long as the dimensionality is properly aligned, after stacking multiple layers of this model, the effective transformation is actually forced to be constrained: where clearly after the first layer the actual feed-forward transformation becomes W W , which is necessarily positive semi-definite (PSD). Hence unconstrained layer weights are not actually possible within this paradigm.
Optimization Induced Attention Mechanisms There is also limited work that attempts to interpret attention mechanisms from the unfolded optimization perspective [14,37]. For example, the energy function was proposed in [37], with subsequent updating using the concave convex procedure [51] leading to iterations of the form y (t+1) = S softmax(βSy (t) ), where S ∈ R n×n corresponds with the attention keys while y maps to the attention query, and β is a constant scalar. Critically though, this attention mechanism is actually cross-attention (using y to attend to S), which does not well-align with the typical self-attention use-case of Transformers. Moreover, this work does not account for the aggregated Transformer feed-forward network module and attendant nonlinearity.
Optimization Induced Graph Neural Networks A variety of graph neural network architectures have also been developed from a similar optimization perspective [7,29,31,34,48,49,52,54]. For example, graph attention mechanisms were derived using the iterative reweighted least squares (IRLS) algorithm in [49]; this result is related to self-attention, which can be viewed as graph attention on fully connected graphs [5]; however, it fails to produce the Transformer softmax term or the combined self-attention/feedforward Transformer stack. Nonetheless, we will later demonstrate how ideas from [49] can be leveraged to derive self-attention with softmax.

Key Challenges Extending to General Transformers
In this paper, we focus on a more complex model than prior work: the Transformer encoder. A Transformer layer [43] is typically composed of two primary components: the self-attention layer and the feed-forward network (FFN) layer. If we simplify the FFN to a single linear transformation with non-linear activation, and ignore the layer-norm and residual connections, 4 one Transformer layer reduces to the basic form With respect to (8) and the limitations of prior work discussed above, we propose four major challenges in formally applying unfolded optimization to the Transformer setting: 1. Token-Level Interactions: As mentioned previously, Transformer models include not only interactions between neurons, but also cross-token interactions, which is reflected by the self-attention mechanism considered to be the sine qua non of Transformers. And thus far, no prior work has derived Transformer-style self-attention, instead either relaxing to cross-attention [37], or failing to produce the ubiquitous softmax operator [49].
2. Heterogeneous Layer Types: Each Transformer encoder layer is composed of two totally different components, a self-attention layer and a feedforward network layer (FFN). We refer to this structure as "heterogeneous layer types," where each component has its own parameters and can be viewed as distinct unfolded optimization processes. However, it remains unknown whether or not it is possible to combine the corresponding energies of these two components to obtain a unified objective with any kind of convergence guarantees (even approximately) during the aggregated forward pass.
3. Non-Linear Activations: Many activation functions used in common neural network architectures can be understood from the perspective of proximal operators [27,49], at least within the isolated context of some relatively simple feedfoward schema. However, when integrated within the heterogeneous Transformer layer types mentioned above, the optimization process becomes rather complex, and it is unknown if any convergence properties still hold after the inclusion of a proximal step.
4. Asymmetry of Weights: As alluded to in the discussion of (5), the scope of most work addressing feed-forward networks is actually limited to models with symmetric (or even tighter, PSD) weight transformations [1,13,47], which restricts the resulting universality. Additionally, related models from the graph neural network literature are generally predicated on undirected graphs [48,49], where the graph propagation is also symmetric. Consequently, although to our knowledge not addressed/discussed previously, it remains unknown how to construct energy functions for more general asymmetric transformations. In this paper, we propose techniques to at least partially solve Challenges 1, 2 and 3, while for 4, we defer preliminary discussion to the supplementary, leaving formal investigation as a future direction.

A New Derivation of Transformer Self-Attention
We now tackle Challenge 1 by constructing an energy function whose iterative optimization steps match Transformer-style softmax self-attention on fully connected tokens as is customary.

Basic Softmax Self-Attention via Unfolded Optimization Steps
Consider the energy function where y i ∈ R d×1 is the i-th row of matrix Y ∈ R n×d , ρ : R + → R is a concave non-decreasing function, and R : R n×d → R is a convex function. Interestingly, although E 1 is not necessarily convex because of the non-convexity of ρ, under specific choices of ρ and R, it can be optimized through a softmax-like structure as follows: let Y (t) represent any fixed value for Y . Then the update rule It is worth mentioning that, although perhaps not obvious, the update step (10) can be generated via a majorization-minimization (MM) algorithm [42], where the majorization step generates a convex upper-bound, and the minimization step descends along the gradient of the upper-bound. Therefore, the core of this update is essentially a gradient step on a convex function, which will be relevant to the discussion in subsequent sections; see the proof and further details in the supplementary. Remark 3.2. Several notable generalizations of Theorem 3.1 are possible. First, although in Theorem 3.1 we adopt a specific form for ρ and R to recover the softmax operator, with other selections of these functions the corresponding unfolded optimization algorithm can generate different/novel types of attention mechanisms. Secondly, obtaining (10) relies on a particular choice of gradient step size; however, for broader choices the resulting convergent updates interestingly lead to residual connections as a natural byproduct of the optimization trajectory (see supplementary). And finally, Theorem 3.1 can be easily modified to accommodate situations whereby full Transformer connectivity is constrained by some graph structure as in [9,17] (again, see supplementary for details).

Extension to Include Trainable Parameters
After aggregating into matrix form, we have thus far shown that the iteration will reduce (or leave unchanged) the energy from (9), where softmax β (y) i = β i exp{yi} j β j exp{yj } denotes a softmax operator with reweighting coefficient vector β. If β i is independent of i, i.e., y (t) i is constant, then this reweighted softmax is equivalent to the canonical softmax used in Transformers.

Now consider the reparameterization
where W s a = W a W a and we have adopted the superscript 's' to indicate that this matrix is symmetric. Collectively then, the results of this section directly address Challenge 1, closely reproducing softmaxstyled self-attention both with and without trainable parameters (the primary lingering limitation of symmetric weights being related to Challenge 4).

Combining Transformer Components via Alternating Inexact Minimization
We next address Challenge 2 and the heterogeneous Transformer layer types. With this goal in mind, we first introduce a general optimization scenario. Specifically, we ask the following question: Given two (convex) objectives f (y) and g(y), under what conditions, and/or to what extent, will alternatively taking separate gradient steps w.r.t. f and g optimize the aggregated function f + g?
We refer to this optimization strategy as alternating inexact minimization (AIM), and we will shortly demonstrate conditions under which it converges to a ball with finite radius containing the optimal point of f + g. Later we discuss how these results contribute towards the resolution of Challenge 2.
We emphasize upfront that AIM as so-defined is quite different from what is commonly referred to as alternating minimization in the literature [18,30]. The latter refers to scenarios whereby a unified objective function with multiple variables is minimized in an alternating fashion over one variable at a time with the others fixed, such that descent can be trivially achieved. In contrast, our AIM scenario involves multiple objective terms with a shared variable, and we alternate minimization over each term in isolation using the same variable, a much more challenging process to analyze.

General Alternating Inexact Minimization
Given two objectives f, g : R d → R, AIM is formalized as Algorithm 1. We will now investigate how Algorithm 1 relates to minimization of the combined objective defined by h(y) = f (y) + g(y).
In the following, we assume that f, g are both Lipschitz smooth and strongly convex [38], with Lipschitz constants L f and L g respectively, and convexity parameter c f and c g respectively. Therefore, h The value of energy E Y (t) with respect to iteration t. Note the y-axis is translated by a constant log E min to make the minimum align with zero.
will also be Lipschitz smooth and strongly convex, with Lipschitz constant L h and convexity c h . We denote the optimal points of f, g, and h as y * f , y * g , and y * h respectively. Combining the two steps of Algorithm 1 gives while a canonical gradient descent step on objective h is Comparing the two update rules it is apparent that Algorithm 1 can be viewed as a noisy gradient descent step with step size α 2 and noise factor Noisy gradient descent is a well-studied problem in certain contexts [4,33]. However, in our specific scenario we require a novel, different (actually tighter) bound than one can get from simply applying existing results.
Specifically, we demonstrate that when δ y (t) = ∆ t / ∇h y (t) is bounded, (13) is guaranteed to descend the objective h within certain region as follows: and y (t+1) are related by (13), with δ y (t) ≤ C and C = α2 α2−α1+α1α2Lg . Then h y (t+1) ≤ h y (t) . Remark 4.2. Although the definition of C may appear rather complicated, when one of α 1 , α 2 , L f , or L g is sufficiently small, C behaves as Ω [α 1 α 2 L f L g ] −1 .
We further consider how to interpret the constraint δ(y) ≤ C and the region within which (13) can optimize h as follows: Combining Theorem 4.1 and Lemma 4.3, we can therefore conclude that when is the whole space excluding a ball centered on y * h . Figure 1 provides a 2D visualization for each case respectively. Remark 4.4. Our findings above can be summarized as follows: For sufficiently small values of α 1 , α 2 , Algorithm 1 reduces the combined objective h, at least provided that y is a certain distance away from the optimal point y * h .
To illustrate this conclusion, we present a synthetic example with F . Note that here we expand the variable y to a matrix Y and let f and g consist of left and right transformations of Y respectively, with an extra bias term to prevent the degenerate solution (Y * = 0). We randomly set each entry of {S, W, B 1 , B 2 } and execute Algorithm 1 with fixed step sizes and project the trace of Y (t) to a two-dimensional space with PCA for visualization. The trace of Y (t) is displayed in Figure 3 and the combined objective h across iterations in Figure 2. From these plots, the behaviour predicted by our theory can be verified: When Y (t) is a sufficient distance away from Y * h (the energy is relatively high), the optimization trajectory moves closer to Y * h (descending the energy), and after Y (t) is close enough to Y * h (and the h is at a relatively low level), the energy oscillates in a certain range about the optimal solution.

Unfolded Optimization of Heterogeneous Layer Type Models
We now return to Transformer-specific models.
analogous to h from the previous section) and alternatively perform the following AIM updates: where provides an additional linear transformation after the softmax.
From Section 3, we know that (15) is optimizing E 1 using essentially a row-wise gradient descent step after majorization. Likewise, (16) is also a gradient step with step size α 2 . Therefore, the combined rules fall into the scope of Algorithm 1, and from the analysis in Section 4.1, we can draw the following conclusion: (15) and (16) when Y (t) ∈ S, where S is a ball with finite radius containing Y * E , the optimal point ofÊ, the convex upper bound of E.
Note that after combining (15) and (16), the update aggregated update rule is already very similar to our target from (8), the only differences being the missing non-linearity and symmetric weights (note we use symmetric weight matrix W s a and W s f instead of W a and W f here), which correspond with Challenges 3 and 4 in Section 2.3 respectively; we address the former next.

Non-linear Activations Integrated within the Unfolding Paradigm
Handling non-linear activations within an unfolded optimization setting has been considered previously [15,41,45,49]. However, prior work has largely relied on proximal operators to create non-linear activations paired with simple linear filters, and the analysis does not transfer to the more complex Transformer case under investigation here. While we will restrict our attention to adding ReLU activations in this section, both for simplicity of exposition and because of their ubiquity in Transformer models, our results generalize to broader selections.
Similar to Section 4, we first study a general form of the optimization problem by introducing a ReLU activations into Algorithm 1, formulating this modification as AIM with proximal steps, and then analyzing convergence under appropriate constraints. Later we apply these results to the Transformer.

Algorithm 2
For the t-th iteration, execute

Proximal-Alternating Inexact Minimization
It has already been established [27,41] that ReLU activations can be modeled as a proximal operator via ReLU(y) = arg min z where φ is the indicator function for the first quadrant given by in Section 4, we have demonstrated how the steps of Algorithm 1 together form an inexact gradient descent iteration of the loss h(y) = f (y) + g(y) with noisy term ∆ t defined in (14). Here, with the addition of the proximal step in Algorithm 2, we obtain an inexact version of proximal-gradient descent. In fact, one turn of Algorithm 2 is equivalent to which is an inexact version of canonical proximal-gradient step with noise ∆ t , as compared to the exact version y (t+1) = arg min z P z; y (t) , where P (z; y) is the proximal problem While various forms of inexact proximal gradient descent have been studied in the past [12,16,39,46], existing work either assumes constant noise [12], stochastic noise [46], or decreasing/convergent noise [16,39]. Critically, no prior work that we are aware of applies to our cse where the noise can potentially increase with iterations. Moreover, existing analysis in the literature is primarily concerned with convergence to a fixed point, while in our case, we instead consider entering a specific region formed around certain points.
Moreover, in addition to bounding δ y (t) as in Section 4, we also need to bound the similarity between the current position y (t) and the gradient α 2 ∇h y (t) , which is defined as D(ξ 1 , ξ 2 ) = 1 ξ 1 2 d i=1 min(ξ 2 2,i −ξ 2 1,i , 0). Intuitively, D(ξ 1 ; ξ 2 ) is defined such that each term of the summation is negative, but only close to 0 when ξ 1,i and ξ 2,i are both large. We then have the following: h , D α 2 ∇h y (t) ; y (t) ≥ −κ for any κ ∈ (0, 1), and δ y (t) ≤ C , where C = Intuitively, Theorem 5.1 shows that the region where the descent of h(y)+φ(y) is guaranteed is the intersection of S(C ) with S defined in Lemma 4.3, and the area T (κ) = {y |D(α 2 ∇h(y); y) ≥ −κ }. While the shape of T (κ) remains difficult to specify in general, we note that T (κ) tends to the whole space when α 2 → 0 or κ → 1. For example, we illustrate the area of T (κ) using a 2D synthetic example, where the energy function is h(y) = W y 2 F + y − b 2 , and entries of W and b are randomly generated. See Figure 4 for the visualization using different values κ, which shows that when κ is sufficiently small, T (κ) is nearly the whole space (except a small area around the origin) which ensures the descent of h(y) + φ(y) in most cases.

Embedding the Non-Linearity within Transformer Models
By adding the penalty term to the previous energy, and with E 1 and E 2 obtained from Sections 3.1 and 4.2 respectively, we finally arrive at the total Transformer energy E(Y ) = E 1 (Y ) + E 2 (Y ) + φ(Y ). And the unfolded optimization of E(Y ) falls into the scope of Algorithm 2, such that the former analysis and Theorem 5.1 apply and we may conclude that the aggregated update rule is a descent algorithm of E, except a region with finite measure. Moreover, if we modify the underlying gradient descent algorithm to make the step sizes α 1 and α 2 sufficiently small, the size of the exception region will tend to zero, and the update rule will be equipped with a residual term (see supplementary; similarly for a discussion about β). Besides these issues, the unfolded update (20) only differs from (8) in its reliance on symmetric weights. This corresponds with Challenge 4 that we have not fully solved, although it has been shown that symmetric weights can mimic asymmetric weights if we enlarge the representation dimensions [22,50].

Practical Verification
Although we have rigorously derived convergence criteria whereby Transformer layers descend a well-specified energy function to a region around optimal solutions, the analysis admittedly relies on conditions that would be difficult to formally verify on real-world datasets. However, our results are nonetheless amenable to targeted empirical corroboration, whereby we can check if the proposed energy does in fact descend during the Transformer forward pass on typical real-world benchmarks.
To this end, we implement a Transformer model as described previously, up to known limitations like symmetric weights. We apply this model to two datasets, IMDB [32] and SST2 [40], which are both commonly-used sentiment classification datasets that use Glove-840b-300d [36] as the word embedding. Figures 5 and 6 display the energy of the output of each layer averaged over 200 randomly chosen samples in the test set. From these figures it is clear that even with real-world data, the the Transformer energy we have derived is (on average) monotonically decreasing across layers, matching the predictions of our analysis. Moreover, with 12 layers (which represents a depth range not uncommon in practice) the model has not entered the fluctuation region. Hence although seemingly complicated, the conditions adopted by Lemma 4.3 and Theorem 5.1 are nonetheless likely hold in many practical settings.

Conclusion
While our contributions here have been primarily of a theoretical nature, there are nonetheless takehome messages that may be of practical relevance. First, as our overriding goal was to reproduce Transformer layers as closely as possible using an unfolded optimization perspective, very specific design choices were made in constructing the core underlying energy function. However, in practice we are free to choose alternative energies that can lead to different forms of tailored self-attention that may be be advantageous on an application-specific basis. Additionally, the distribution of attention weights can potentially be better understood or influenced by properties of the energy function that produces them. And finally, especially for regimes with limited data and therefore fewer free model parameters, the unfolding perspective can be used to devise architectures with inductive biases aligned with downstream tasks to help compensate for less model flexibility.
We conclude by noting that the techniques introduced in this paper can in some sense be viewed as universal tools for constructing a broader family of unfolding architectures. In fact, many/most deep learning models, including those with feed-forward structure [28], token-level interactions [44], residual connections [20], and heterogeneous layer types, can be reinterpreted through the lens of unfolded optimization using the framework we have introduced, at least up to some potential constraints like symmetric weights (which may be handled in other ways).

A Theory Background
This section provides background and references regarding some of the concepts and techniques used for proving our main results.

A.1 Lipschitz Conditions and Strong Convexity
For a function f : we say f satisfies the Lipschitz condition with Lipschitz constant L or is simply L-Lipschitz. If the gradient of a function is L-Lipschitz, we say it is Lipschitz smooth with Lipschitz constant L.
we say it is (strongly) convex with convexity c, or simply c-strongly convex.
There are two important inequalities for Lipschitz smooth and strongly convex functions: If f is c-strongly convex, then The proof of Proposition A.1, and more details of these two properties can be found in [4].

A.2 Proximal Operators and Proxinal Gradient Descent
The proximal problem of a function φ with parameter λ is defined as and the mapping from y to the minimal point of P (·; y), is called the proximal operator of φ: It can be shown that, for an objective function h = f + φ, where f is smooth, alternatively performing gradient descent on f and proximal projection of φ, i.e.
is a descent algorithm of h, which is often referred to as proximal gradient descent [35].

A.3 Subgradient
For a (not necessarily smooth) convex function f : R d → R, it can be proven that the following set is always not empty [8]: We call ∂f (y) the subdifferential of f at point y, and the elements of ∂f (y) the subgradients of ∂f at point y.
Subdifferential satisfies linearity, which means where the summation of two sets is defined as the elementary summation Although subdifferentials are sets, based on the linearity, it is easier to use the equation symbol = instead of ∈ to denote subgradient, e.g. ∂f (y) = a means a ∈ ∂f (y).
One very useful property of the subgradient is that it can indicate the global minimum of a convex function : which is quite straight forward from the definition of subgradient. More detailed introduction and discussion of subgradients can be found in [6,8].

B Proofs, Discussion, and Extensions of Section 3 Results
We now provide a complement of Section 3 of the main paper, including the proofs and discussion. For Theorem 3.1, we prove a more generalized version, taking into account of the structure of attention and the step size, as mentioned in Remark 3.2.
To achieve this, we extend the techniques used in [49] and show how to construct an energy function whose iterative optimization steps match Transformer-style self-attention on arbitrary graph structures -including fully connected graphs.
Note that this section provides a more general framework that assumes a certain graph structure exists in token-level interactions, which covers some attention sparsification works like [9,17]. And when the graph is fully connected, we obtain a normal Transformer structure as analyzed in the main paper.

B.1 General Unfolded Optimization Results for Creating Self-Attention on Graphs
Suppose G = (V, E) is an undirected graph, whose adjacency matrix is A ∈ R n×n , incidence-matrix is B ∈ R m×n , degree matrix is D = diag (A) 1, and Laplacian matrix is where n = |V| and m = |E|. We use a tuple (u, v) to represent an edge in the graph that connects node u and v. Furthermore, let there be an index for each edge, suppose the ith edge connects fr(i) and to(i) and the index of edge that connecting (u, v) is ind(u, v). Then, the incidence matrix B can be defined as Consider a generalized version of the energy function (9): where y u ∈ R d×1 is the u-th row of matrix Y ∈ R n×d , ρ : R + → R is a concave non-decreasing function and R : R n×d → R is a convex function.

Algorithm 3
Execute the following two assignment operations in arbitrary order: where Proposition B.1 (Lemma 3.2 in [49]). Let Y (t) be the input of Algorithm 3 and Y (t+1) be the output, In Algorithm 3, Step 1 converts the non-linear function ρ to edge weights in graphs, which correspond to attention weight matrices. If in Step 2 we use gradient descent, then the update rule is where the attention weights are injected into coefficients 1 − B ΓB. However, this naive framework does not include softmax used in Transformers. In the following subsections we show that with a special choice of ρ and a reweighting of each row of Y , we get exactly self-attention updates with softmax.

For
Step 2, consider using gradient descent with Jacobi preconditioning, from which we arrive at the following Theorem: Theorem B.2. Consider updatingẼ 1 using a gradient step with step size α and Jacobi preconditioner D −(t) : where and α ≤ 1, it follows thatẼ 1 Y (t+1) ≤Ẽ 1 Y (t) . And the update rule in (37) can be written as Notice that it requires α ≤ L −(t) to guarantee (37) is a descent step ofẼ 1 , where L (t) is the Lipschitz constant of the gradient of D −(t)Ẽ 1 Y, Γ (t) [4]. And by the definition of D (t) , L (t) is always 1. Therefore, α ≤ 1 always satisfies the condition.

B.3 Discussion
From the analysis above, we have the following observations. Firstly, if we set α to other values, say α = 1 2 , then the residual term in (10) is maintained, which corresponds with the common use of skip connections in GNNs. Secondly, by assuming arbitrary graphs G (as opposed to merely the fully connected case described in the main paper), this framework can naturally handle attention mechanisms with special structure as in [9,17].

C Proofs of Section 4 Results
In this section, we prove the propositions in Section 4.

C.1 The Proof of Theorem 4.1
Firstly, from the Lipschitz smooth assumption of the objective functions, we can bound the norm of the noise term: In the following we write ∇h for short to refer to ∇h y (t) ; similarly δ for δ y (t) and ∆ for ∆ t .

C.2 The Proof of Lemma 4.3
We only need to notice that ∇f y * f = ∇h (y * h ) = 0, then by Lipschitz smoothness and convexity of the objectives, Therefore, when y ∈ S(C ),

D Proofs of Section 5 Results
For simplicity, in the following we use y to denote y (t+1) and x to denote y (t) this section. We also use x * to denote the optimal solution of the proximal problem (19): Recall that y = arg min z First we shall prove the bound of the subgradient of proximal problem (19) at the point of y: Proof. Since y is the optimal point of (59), we have ∂ ∂y which gives Therefore, the subgradient of (19) at y is Then, we shall show the descent of P (y; x) to P (x; x) can be bounded by the distance between x and x * . Lemma D.2.
Proof. Let c P denote the convexity of proximal problem P (z; x). And note since x * is the optimal point of P (z; x), we have ∂P (x * ) = 0. Then by convexity we have and (By (45) (71) Combining (66) and (71) we prove the lemma.
Given Lemma D.2, we only need to bound x − x * . The following lemma shows that, x − x * is related to as stated in the main paper. Lemma D.3. Given D (α 2 ∇h (x) ; x) ≥ −κ, we have Therefore Next, we shall show that P (z; x) is actually an upper bound of f (z) + φ(z): 2λ ∇h(x) 2 , then by the assumed Lipschitz condition we have = P (z; x) + η(x).
And it is also straightforward to verify that It is worth noting that, although there's a parameter λ in the definition of P , when φ is defined as , which we used in the main paper to derive ReLU, the proximal operator prox φ is independent of λ, which means we can always select a λ small enough to satisfy λ ≤ α 2 , which is required by Lemma D.4.
Lemma D.4 provides us with a majorization minimization perspective, which allow us to bound h(z) + φ(z) by bounding P (y; x). And Lemma D.2 provides a way to bound the descent of P (y; x). Now we are ready to prove Theorem 5.1.

E Further Discussion
In this section, we provide further discussion related to symmetric weights, layer-dependent weights, and layer normalization as mentioned in main paper.

E.1 Asymmetric Weights
It initially seems hard or impossible to directly interpret an asymmetric transformation (i.e. f (y) = W y where W is asymmetric) as a gradient (See Lemma 5.3 in [50]). However, non-PSD weights are possible to be viewed as a gradient since ∂y W y ∂y = W + W y, where W + W is symmetric but not necessarily PSD. Moreover, if the activation function is linear or meets certain criteria, it is also possible to use symmetric weights [50].
Furthermore, there is other work showing that models with symmetric weights are also universal approximators [22], which suggests that using symmetric weights might not adversely affect the expressivity of the model provided the hidden dimension can be increased sufficiently.

E.2 Layer-dependent Weights
Note that since in the unfolding framework every step is the optimization process of a certain energy function, the model must exhibit a characteristic of recursive models like in (4), where the same weight matrix is shared at each layer. However, as has been discussed in various contexts, if the number of model layers is finite, it is possible to use shared weights to "simulate" layer-dependent weights [3,47,50]. Moreover, from the construction in [50], when the number of model layers is finite, the weights can even be asymmetric.

E.3 LayerNorm
In the main paper, the model is actually a simplified version of the Transformer since LayerNorm is missing. Here we note that it is possible to combine LayerNorm in our proposed unfolding model. Apart from parameters, LayerNorm can be viewed as following process: 1. Translation: i , 2. Rescaling: y (t+1) = u (t) / u (t) ; which can both be interpreted as gradient steps: 1. Translation: u (t) = y (t) − ∂ Therefore, using the techniques in Section 4, we can insert LayerNorm in an unfolding model by adding extra energy terms. Moreover, with LayerNorm applied, the reweighted softmax in (40) becomes normal softmax.