Transformers from an Optimization Perspective
Authors: Yongyi Yang
Zengfeng Huang
David Wipf
Venue: NeurIPS 2022 Accept (2022)
Keywords: Transformers
self-attention
unfolded optimization
energy function minimization

Abstract: 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.  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.

Transformers from an Optimization Perspective
Yongyi Yang' University of Michigan yongyi @umi ch. edu
Zengfeng Huang
David Wipf Amazon Web Services
Fudan University Iuangzf@fudes edu.
pnrrdwiip@@gaaae com com
Abstract
Deep learning models such as the Transformer are often constructed by heuristics and experience. To provide complementary foundation, in this work we study the following problem: Is 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 view Transformers the unfolding of 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-altention mechanisms like the Transformer. To this end, first outline several major obstacles before providing companion techniques to at least partially address them, demonstrating for the first time close association between energy function minimization and deep layers with self-attention. This interpretation contributes our intuition and understanding of Transformers, while potentially laying the ground-work for new model designs.
Introduction
Although deep learning has achieved extensive practical success 221 model architectures often motivated by heuristics and at times remain difficult to understand. As such, significant effort has recently be devoted various ways of analyzing and reinterpreting popular deep network structures. Of particular note for our purposes here are assorted attempts to create a one-to-one association between deep model layers and the unfolded iterations of some optimization process designed to minimize an interpretable energy function [10 viewed as computing, at least approximately, representation with minimal energy, where model parameters dictate the form of this energy and can be trained via the backward pass for downstream goal of interest, e.g., classification regression, etc.
35]. In this way the forward pass can be
This so-called unfolded optimization perspective can provide insights into the nature of the inductive biases introduced by different architectures, while potentially serving guide for bespoke design choices instantiated through properties of the underlying energy function involved. fact, modifi- cations of the latter can lead to predictable model behaviors, which have in the past been tailored to incorporate useful regularization factors [53]. avoid subtle degenerate solutions [49] or design effective new models altogether [O].
Despite these successes, the majority of prior work this area has either addressed relatively simple network structures such as MIPs 47 or else restricted consideration to single model component in isolation from a larger system 35 Consequently, complex models such as the widely-used Transformer
43 have been mostly ignored, in large part because of the difficulty involved in simultaneously mapping both the self-attention and feed-forward network (FFN) modules to an integrated optimization process that minimizes single, unified energy function.
"Work completed during an internship the Shanghai Lab.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

We take significant steps towards extending the unfolded optimization perspective to full Trans- former stack via the following contributions:
After providing background details and existing examples of unfolded optimization, Section 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 deerives an energy function that closely reproduces self-attention per Chal- lenge (i), Section rigorously derives convergence results which demonstrate how to approximately handle Challenge (ii) via novel form of alternating minimization, and Section 5 provides explicit details for solving Challenge (iii) via proximal methods; consid- eration of Challenge (iv) deferred to the supplementary along with additional supporting analyses.
The above contributions culminate in an energy function that is minimized by Transformer layers when certain technical conditions hold. To provide further motivation, Section 6 empirically demonstrates that indeed this energy is minimized with real-world data even when technical conditions 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 the unfolded optimization perspective is create a one-to-one correspondence between the layers of a deep model and the iterations of some algorithmic process designed to minimize parameterized, model-specific energy function. In this way, the forward pass can be viewed as computing an approximation to the minimum energy solution. Mcanwhile, assuming each forward iteration 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 plleve-optimizztionn process 77 I0 20). Overall then, unfolded optimization produces minimal energy representations, whereby the specific structure of the energy function is trainable per the relevant downstream application domain.
2.1 Mathematical Formulation
The bilevel optimization process described above can be expressed more formally as:
(w,,,) x) = arg min F(Y; X), Y
W arg min
1 5
Σ
Here S is the size of dataset, & is the i-th loss function and X, E Raxd is the i-th input instance (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 E x as the hidden representation the (ideal) model output (e.g. set of word representations output by a Transformer encoder) that minimizes to lower-level energy functior : nxd + R whose form is dictated by parameters W (e.g., these could also correspond with the trainable parameters in a Transformer model). And f a meta-loss function (c.g., the loss function used to train a Transformer) while & : R + is an arbitrary transformation that converts hidden representations to the model output, which can also potentially have its own learnable parameters. Assuming a x)/OW and ae/ Df/0Y*(W; X) 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.
the model, where Y is
Note that hereinafter if the energy function does not include inter-token interactions, we view it as function of vectors, and we use lower-case characters y and æ to represent an arbitrary row of and
Note that although we frequently the term "energy", the concept is different from what has been called 'energy-based learning" in the past, where the energy is defined on the joint input and output spaces 22 23 in energy is just function for explaining the forward pass.

---

X for simplicity. Also, when clear from context, we omit writing the parameters of function, for example we may write E to reference E(Y; X.).
2.2 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-Inducee Feed- Ferd-Forward Networks wide variety of prior work has investigated feed-
forward structures from unfolded optimization perspective [10] 12 16
[8] 41 47 Here we take
closer look at [47 which is based on proximal methods. The basic energy function in [47]
E(y) = 1 (y) (f(x), W Ty)
|y1
(3)
where 1 is vector with all entries being 1, o² is the convex conjugate 336 of a+ So (r))dr, is activation function and (a=) is some transformation of input presentation æ. The proximal operator 8 of is
((++)) proxg (xx(0) -WT (Wy())((x))
+
(4)
which is analogous the feed-forward layer of MLP. Note that although W 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:
y(++2) WT a [WWTo (wy + f(x)) + f(x))
(5)
where clearly after the first layer the actual feed-forward transformation becomes WWT, which necessarily positive semi-definite (PSD). Hence unconstrained layer weights are not actually possible within this paradigm.
Optimitation-Induuce Attention Mechanisms There is also limited work that attempts to interpret attention mechanisms from the unfolded optimization perspective 35). For example, the energy function
d
Ely) = -3--1 log
+
2 |yl?
(6)
was proposed in [35], with subsequent updating using the concave convex procedure [5] leading to iterations of the form
y (1+1) = S softmax(3SS(1))
(7)
where S E _ax= corresponds with the attention keys while y maps to the attention query, and is a constant scalar. Critically though, this attention mechanism 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.
OptimizationnIInuaee Induced Graph Neural Networks A variety of graph neural network architectures
have also been developed from similar optimization perspective [ 26 29
32
48
49
52
54 For
example, graph attention mechanisms were derived using the iterative reweighted least squares (IRLS) algorithm in 49 While this result is related to self-attention, which can be viewed as graph attention on fully connected graphs 4) fails to produce the Transformer softmax term or the combined tuffaattentiiooffeefforaaat Transformer stack. Nonetheless, we will later demonstrate how ideas from (49 can be leveraged to derive self-attention with softmax.
2.3 Key Challenges Extending to General Transformers
In this paper, we focus on a more complex model than prior work, namely, the Transformer encoder. Transformer layer 43 typically composed of two primary components: the self-attention layer and the forwaforward network (FFN) layer. If we simplify the FFN to single linear transformation
Herr o² is scalar function, but for simplicity, we sometimes apply scalar function to vector or matrix entry-wise fashion for convenience.
3

---

with non-linear activation, and ignore the layer-norm and residual connections + layer reduces the basic form
one Transformer
y(++)) = RcLU softmax (y (r(lo)))))) rro),,)
(8)
where Wa and W represent self-attention and FFN layer weight matrices respectively. Regarding (8) and the limitations of prior work discussed above, we propose four major challenges in formally applying unfolded optimization to the Transformer setting:
Token-Level Interactions: As mentioned previously, Transformer models include not only a feed-forward process but also cross-token interactions, which is reflected by the self-attention mechanism considered to be the sine que non Transformers. And thus far, no prior work has derived Transformer-styy
self-attention, instead either relaxing to cross-attention [35]. or failing to produce the ubiquitous softmax operator [49].
2. Heterogeneous Layer Types: Each Transformer encoder layer is composed of two totally different components: self-attention and a FFN. We to this structure as "heterogeneous layer type," where each component has its own parameters and can be viewed unfolded optimization process. However, remains unknown whether or not is possible to combine the corresponding energies of these two components to obtain unified objective with any kind of convergence guarantees (even approximately) during the aggregated forward pass.
distinet
3. Non-Linear Activations: Many activation functions used in common neural network ar- chitectures can be understood from the perspective of proximal operators (24 49 least within the isolated context of some relatively simple feedforward schema. However, when integrated within the heterogeneous Transformer layer types mentioned above, the optimiza- tion process becomes rather complex, and it is unknown if any convergence properties still hold after the inclusion of 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 [ 10 47). which restricts the resulting universality. Additionally, related models from the graph neural network literature are generally predicated undirected graphs 48 49), where the graph propagation is also symmetric. Consequently, although to our knowledge not addressed/diseussed previously, remains unknown how to construct energy functions for more general asymmetric transformations.
In this paper, we propose techniques to least partially solve Challenges 1 and while for 4 defer preliminary discussion to the supplementary, leaving formal investigation as a future direction.
3 A New Derivation of Transformer Self-Attention
We now tackle Challenge by constructing energy function whose iterative optimization steps match Transformer-style softmax self-attention on fully connected tokens as is customary.
3.1 Basic Softmax Self-Attention via Unfolded Optimization Steps
Consider the energy function
E
1 ly. -,11²) + R(Y),
(9)
where y, Rdx1 is the i-th row of matrix Y E Rxad p R IR is a concave non-decreasing function, and R Raxd - IR is a convex function. Interestingly, although E, is not necessarily convex because of the non-convexity p, under specific choices and R, it can be optimized through a softmax-like structure as follows:
*Note that we are ignoring layer-norms and residual connections for simplicity of exposition; however, the analysis we present herein can naturally be generalized to accommodate layer-norm, and also certain forms of residual connections with additional assumptions: see supplementary for details.

---

Theorem 3.1. Assume p ((2) = exp {-2], R(Y) = 2 2 and 3, exp { { yx 13 and
let (()) represent any fixed value for Y. Then the update rule
Ej exp { (1)T (1) (t)
y1 ((+1))
Vi,
(10)
Ej exp ( (7)
satisfies Ei (y(++)) S Ei (Y with equality iff ( a stationary point of Ey.
It is worth mentioning that, although perhaps not obvious, the update step (10) can be generated via a
(MM) algorithm 42). where the majorization step produces a convex poper- and the minimization step descends along the gradient of the upper-bound. Therefore, the core of this update is essentially 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 [31 are possible. First, although in Theorem 3.I w. adopt a specific form for p 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 (TO) relies on 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 constrained by some graph structure as 6 [4] (again, see supplementary for details).
3.2 Extension to Include Trainable Parameters
After aggregating into matrix form, we have thus far shown that the iteration
((+1) softmaxx (y (1) (()) y
(1I)
will reduce (or leave unchanged) the energy from
9 where softmax,(y):
Σ, exply, denotes
9 exply.)
a softmax operator with reweighting coefficient vector B. If (3 independent of i, i.e., y
(2)
is
constant (which is enforceable via layer normalization that can also be included within our framework as mentioned previously), then this reweighted softmax is equivalent to the canonical softmax used in Transformers.
Now consider the reparameterization Y ZW where Wa Rddd is an invertible matrix. It follows that (+++) W. softmax,3 (Z W, W. Z T) Z(,,,, a, leading to the revised update
z(++1) softmaxg
2 W22(0T) z(1)
(12)
where W = W.W and we have adopted the superscript to indicate that this matrix is symmetric. Collectively then, the results of this section directly address Challenge I closely reproducing softmax- styled self-attention both with and without trainable parameters (the primary lingering limitation of symmetric weights being relegated to Challenge
4 Combining Transformer Components via Alternating Inexact Minimization
We next address Challenge and the heterogeneous Transformer layer types. With this goal in mind, first introduce general optimization scenario. Specifically, we ask the following question:
Given two (convex) objectives f(y) and gly), under what conditions, and/or 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 alternating inexact minimization (AIM), and we will shortly demonstrate conditions under which converges to a ball with finite radius containing the optimal point of 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 commonly referred to as alternating minimization in the literature [15, 28 The latter refers to scenarios whereby unified objective function with multiple variables minimized in an alternating fashion over one variable at 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.
5

---

2.5
log 1ggE
2.0
E(yi0)
1.0
0.5
125
175 200
Figure 2: The value of energy E (Y) with respect to iteration t. Note the y-axis is trans- lated by constant log Emin to make the min- imum align with zero.
Figure I: A two-dimensional illustration of ((6) with different values of '6 Left:
e 6 = 0.7. Right: o by
1.5. The blue area is S(6).
4.1 General Alternating Inexact Minimization
Given two objectives f.g: RRd + R, AIM is formalized as Algorithm where and ag constitute step sizes. We will now investigate how Algorithm I objective defined by h(y)= = f(y) + g(y).
relates minimization of the combined
Algorithm
For the 1-th iteration, execute
(1) y(t) ai Vf(y('));
y(++1) (()) 02 ()(
In the following, we assume that f. g are both Lipschitz smooth and strongly convex [36], with Lips- chitz constants Ly and Lg respectively, and convexity parameter сf and G respectively. Therefore, h will also Lipschitz smooth and strongly convex, with Lipschitz constant L,, and convexity e.. We denote the optimal points of f.g. and h yy, Y and y respectively.
Combining the two steps of Algorithm 1 gives
y ((+1) y an V (5(0)
V2 y(() on VS (y()))]
(13)
while a canonical gradient descent step on objective h is ((+2)) -2 V (xx()) + Vy
(}]
Comparing the two update rules is apparent that Algorithm can be viewed as a noisy gradient descent step with step size a2 and noise factor
A = Vh (1)
an V
(y(1))
Vg -n V (x
(14)
Noisy gradient descent is a well-studied problem in certain contexts 3 31 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 8 (y) = | Vh (y(9) teed to descend the objective h as follows:
bounded, 13
is guaran-
Theorem 4.1. When 01 02 LT suppose ( and andy(++1) (++ are related by 13 with 8 (y(1) & and 6 02 +0-00 Then (y(++1)) S (y())).
02
Remark 4.2. Although the definition of 6 may appear rather complicated, when one of a1, 02 Ly-
or Lg is sufficiently small, & behaves as Ω ([a10يL Lyl-))
We further consider how to interpret the constraint 8(y) < & and the region within which 13 can optimize follows:
Lemma 4.3. Let S(6) -(8)=
yi vill
Cn 36' Ly J. When (1) E S(6), 5(y) < 8.
Combining Theorem 4.1 and Lemma 4.3 we can therefore conclude that when y()) E S(6). (y(++))) (y Note that the boundary of S(66) given Lemma 4.3 is called an Apollonian circle 37 When & ما, S(E) is ball centered on y", and when & LL S(8)) is the whole space excluding a ball centered on y Figure I provides 2D visualization for each case respectively.

---

Additionally, note that the same analysis can be done by switching the role of f and g (since the two processes are alternating) and further restrict the excepted region, although here we omit this for simplicity.
Remark 4.4. Our findings above can be summarized as follows: For sufficiently small values of a1, 022, Algorithm Ireduces predues the combined objective h, least provided that y is a certain distance away from the optimal point yh-
To illustrate this conclusion, we present a synthetic example with f(Y) | |Y Bill' and g = ||YWI] + ||Y B2ll3. Note that here we expand the variable y to a matrix Y and let and g consist of left and right transformation, of Y respectively, with an extra bias term to prevent the degenerate solution (Y = 0). We randomly set each entry of (S, W. Bi, B2} and execute Algorithm
1 with fixed step sizes and project the trace of y to a two-dimensional space with PCA for visualization. The trace of (()) 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 (() is a sufficient distance away from Yr, (the energy is relatively high), the optimization trajectory moves closer to Y (descending the energy), and after (+) is close enough to Y₁, (and the h is at relatively low level), the energy oscillates in certain range about the optimal solution.
60
76-
-20
100 325 150
375
Figure 3:
The trace of (() projected to a plane. The figure on the right displays a zoomed local region around Y taken from the figure the left.
4.2 Unfolded Optimization of Heterogeneous Layer Type Models
We now return to Transformer-specife models. Let Ey = Ir (VW,YT)+2IVII and consider the combined energy E Ei (Y.,)) + Ez(Y)) (analogous to l from the previous section) and alternatively perform the following AIM updates:
U softmax,
(oy) y(u;
(15)
0E2
(++ ((t)
022
= W,
(16)
OY
YY-W)
W,+W7 provides an additional linear transformation after the softmax.
7 where W = - - 022
2
From Section 3] we know that 15)) is optimizing Eg using essentially row-wise gradient descent
step after majorization. Likewise, (16) is also a gradient step with step size 02 Therefore, the combined rules fall into the scope of Algorithm [] and from the analysis in Section 4.1 we can draw the following conclusion:
Corollary 4.5. (1Y(+ is computed via (1) and (6) with input ( then E (Y(+))) S E (Y(1)
when (t) & S, where S is a ball with finite radius containing Y the optimal point of É, the convex upper bound of E.
Note that after combining 11 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 use symmetric weight matrix We and Wi instead of W, and Wy here), which correspond with Challenges 3 and in Section 2.3 respectively; we address the former next.
7

---

5 Non-linear Activations Integrated within the Unfolding Paradigm
Handling non- activations within an unfolded optimization setting has been considered pre- 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 exposition and because of their ubiquity in Transformer models, our results generalize to broader selections.
viously [12
41 45 49,
Similar to Section 4 we first study general form of the optimization problem by introducing a ReLU activations into Algorithm[ formulating this modification as AIM with proximal steps, and then analyzing convergence under appropriate constraints. Later we apply these results to the Transformer.
5.1 Proximal-Alternatinn Inexact Minimization
Algorithm 2
For the 1-th iteration, execute
(1) = y - ai V. (y(");
v(t) = ( 02 Vg (u( y((+1) = ReLU (v(").
It has already been established 24 that ReLU activations can be modeled a proximal operator via
RcLU(y) arg min _ || y1 Φ62),
(17)
where Φ is the indicator function for the first quadrant given by x(z) = +00 if <
Moreover,
0
if 2 0"
in Section 4 we have demonstrated how the steps of Algorithm [Tagetheerr form an inexact gradient descent iteration of the loss h(y) fy) + g(y) with noisy term A, 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 i equivalent
y ((++)) arg min
(()) t2 ( =22. 2 d(z),
(y -
(18)
2X
which is an inexact version of canonical proximal-gradient step with noise Dc, as compared to the exact version y(++1) arg ming P
y())()). where P(z; y) is the proximal problem
P 2: (2)
2X 2 (0) + az V2 )I + d(z).
(19)
While various forms of inexact proximal gradient descent have been studied in the past 9 1
38
existing work either assumes constant noise [9]. stochastic noise [46], or decreasing/converreee
noise [ 38. Critically, no prior work that we are aware of applies our case where the noise can potentially increase with iterations. Moreover, existing analysis in the literature is primarily concerned with convergence to fixed point, while in our scenario, we instead consider entering a specific region formed around certain points.
In addition to bounding 8 (y)))) as in Section we also need to bound the similarity be- tween the current position y(t) and the gradient 0: Vh (y(1)), which is defined as ()), = E Ei min(E2, E. 0). Intuitively, D(E E2) is defined such that each term of the summa- tion is negative, but only close to when E1.s and E2, are both large. We then have the following:
Theorem 5.1. Ifa1= S02 < I D (a2Vh (y(()) (t) 2 -k for any E (0, 1), and 5 (x))) S we have that h (y(++)) + & (y(++)) < h (y + Φ (y()).
aacpAvI-x 8", where 8" V2(02 72(02-01 +1100al))
Intuitively, Theorem 5.1 shows that the region where the descent of hv) + ( is guaranteed is the in- tersection of S((") with S defined in Lemma 4.3 and the area ((x) = {y 1D V)(y);y) 2 - While the shape of T ((x) remains difficult to specify in general, we note that T (x) tends to the whole space when a2 + 0 or k + 1. For example, we illustrate the area of T(x) using a 2D synthetic example, where the energy function is h(y) = ||Wyl]3 + Ily 1-b22 and entries of W and b are randomly generated. See Figure 4 for the visualization using different values which shows that when x is sufficiently small, T(k) nearly the whole space (except a small area around the origin) which ensures the descent of h(y) + Φ(y) in most cases.
8

---

02 6=04
04 0.
Figure 4: T(x) region (blue area) of 2D synthetic example with different values of is
5.2 Embedding the Non-Linearity within Transformer Models
By adding the penalty term to the previous energy, and with E1 and Eg obtained from Sections 3.1 and
4.2 respectively, we finally arrive at the total Transformer energy E(Y) = E, (Y) ) + EY()) + +
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
Y (++1) ReLU softmaxx (r("):)(07) Wi]
(20)
descent algorithm of E, except region with finite measure. Moreover, we modify the underlying gradient descent algorithm to make the step sizes and a2 sufficiently small, the size of the exception region will tend zero, and the update rule will be equipped with residual (see supplementary; similarly for discussion about B3. Besides these issues, the unfolded update (20) only differs from ( in its reliance on symmetric weights. This corresponds with Challenge that we have not fully solved, although it has been shown that symmetric weights can mimic asymmetric
weights if we enlarge the representation dimensions [9 50
Practical Verification
Although have rigorously derived convergence criteria whereby Transformer layers descend well-specified energy function to region around optimal solutions, the analysis admittedly relies on conditions that would 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 benchmarks.
To this end, we implement a Transformer model as deseribed previously, up to known limitations like symmetric weights. We apply this model two benchmarks, IMDB 300 and SST2 are both commonly-used sentiment classification datasets that rely on Glove-840b-300d
40 which
33 as the
word embedding. Figures [nd 6 the energy of the output of each layer of Transformer (as defined in (8) averaged over 200 randomly chosen samples in the test set. Figure 5 juses randomly initialized weights while Figure involvess weights trained for 2000 steps with SGD and learning rate 0.01. Additionally, for the trained model we change the term ||Y 11 in E1 to ||Y X| in order avoid degenerate representations (as can sometimes occur trained Transformers (391), noting that this modification is equally- well covered by our theory and leads to commonly commonlyy-ued form of residual connection the resulting Transformer architecture. See supplementary for details.
08
03
6
3
2
3
3
5
10
21
12
E( E(ayar indext
Blayar index)
Figure 5: Box plot of the energy curves from randomly initialized 12-layer Transformers on IMDB (left) and SST2 (right) datasets; results are averaged over samples.
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

---

3.5
1.0
1
(ayer Indexi
[layer Indes)
Figure 6: Box plots of the energy curves from trained 12-layer Transformers on IMDB (leff) and SST2 (right) datasets; results are averaged over samples.
analysis. Moreover, with 12 layers, which represents a depth range not uncommon in practice (e.g.. BERT), the model has not entered the fluctuation region. Moreover, this observation holds even with trained Transformer models composed from many/most of the typical components used in practice, namely, layers with self-attention, a linear transform followed by nonlinearity, and residual connections, etc. 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 take- home messages that may be of practical relevance. First, 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-specife basis.
As brief/simple representative example of the latter consider the following: The canonical Trans- former uses softmax to normalize the attention coefficients. However, multiple works have questioned the appropriateness of softmax normalization, e.g., [27 34 Under our optimization unfolding perspective, can see that the softmax normalization in the Transformer is derived from a specific choice of p function, namely p(=) -e- (Section 3.1 and Theorem 3.1 But if we choose different p. it would lead to new normalization method other than softmax with interpretable properties. For instance, if p(z) = log(z + 2), then the attention coefficients would behave as (for simplicity we assume all 's have unit norm below):
ΣS O. = 2--2241 - y, A= y
1
-
where azj is the attention coefficient between the i-th and j-th token. Since log(2 + 2) grows slower than - for 2 E [0, 1]. the associated attention formulation would tend to encourage some more dissimilar representations between tokens. This is because in the new energy, large value of lly; yý (which means dissimilar y; and y,) contributes less to the energy compared to before, and thus the optimization process has less motivation to reduce it. Conversely, if choose & that grows faster than -e -= in the stated range (e.g. p(2) = log(2 + 1)), then the derived model would likely be encouraging more similar representations between tokens.
There are other ramifications of this overall framework with practical relevance as well. For example, the actual distribution of attention weights can potentially be better understood or influenced by properties the energy function that produces them. Moreover, 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 broader family of unfolding architectures. In fact, many/most deep learning models, including those with feed-forward structure (25), token-level interactions 44 residual connections [17.. 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).
10

---

References
[1] Brandon Amos and J. Zico Kolter. OptNet: Differentiable optimization as a layer in neural networks. In International Conference on Machine Learning, ICML, Proceedings of Machine Learning Research, 2017.
[2] Brandon Amos, Lei Xu, and J. Zico Kolter. Input convex neural networks. In International Conference Machine Learning, ICMI., 2017.
[3] Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Rev., 60(2):223-33 2018.
[4] Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Velickovic. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. CoRR, abs/2104. 13478, 2021.
[5] Siheng Chen and Yonina C Eldar. Graph signal denoising via unrolling networks. In Interna- tional Conference on Acoustics, Speech and Signal Processing, ICASSP, 2021.
[6] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904. 10509, 2019.
[7] Benoît Colson, Patrice Marcotte, and Gilles Savard. An overview of bilevel optimization. Ann. Oper. Res., 153(1):235-256, 2007.
[8] Patrick L Combettes and Jean Shan-Christophe Pesquet. Proximal splitting methods in signal processing. In Fixed-point algorithms for inverse problems in science and engineering, pages 185-212. Springer, 2011.
[9] Olivier Devolder, François Glineur, and Yurii E. Nesterov. First-order methods of smooth convex optimization with inexact oracle. Math. Program., 146(1-2):37-77 2014.
[10] Jordan Frécon, Gilles Gasso, Massimiliano Pontil, and Saverio Salzo. Bregman neural networks. In International Conference on Machine Learning, ICML, 2022.
[11] Zhengyang Geng, Meng-Hao Guo, Hongxu Chen, Xia Li, Ke Wei, and Zhouchen Lin. Is attention better than matrix decomposition? In International Conference on Learning Represen- tations, ICLR, 2021.
[12] Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In International Conference on Muchine Learning, ICMI., 2010
[13] Bin Gu, De Wang, Zhouyuan Huo, and Heng Huang. Inexact proximal gradient methods for non-convex and non-smooth optimization. In AAAI Conference Artificial Intelligence, 2018.
[14] Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. Star- transformer. In Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-IILT, 2019.
[15] Moritz Hardt. Understanding alternating minimization for matrix completion. In Annual Symposium on Foundations of Computer Science, FOCS, 2014.
[16) Hao He, Bo Xin, Satoshi Ikehata, and David Wipf. From bayesian sparsity to gated recurrent nets. In Annual Conference on Neural Information Processing Systems, NeurIPS, 2017.
[17] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEÉE Conference Computer Vision and Pattern Recognition, CVPR, 2016.
[18] John Hershey, Jonathan Le Roux, and Felix Weninger. Deep unfolding: Model-based inspiration of novel deep architectures. CoRR, abs/1409.2574, 2014.
[19] Shell Xu Hu, Sergey Zagoruyko, and Nikos Komodakis, Exploring weight symmetry in deep neural networks. Comput. Vis. Image Underst., 187, 2019.
[20] Simon Jenni and Paolo Favaro. Deep bilevel learning. In European Conference on Computer Vision, ECCV, 2018.
11

---

[21] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In Annual Conference on Neural Information Processing Systems, NeurIPS, 2012.
[22] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521((553):436 444, 2015.
[23] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and F Huang. A tutorial on energy- learning. Predicting structured data, 1(0), 2006.
[24] Jia Li, Mingqing Xiao. Cong Fang, Yue Dai, Chao Xu, and Zhouchen Lin. Training neural networks by lifted proximal operator machines. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020.
[25] Hanxiao Liu, Zihang Dai, David R. So, and Quoe V. Le. Pay attention to mlps. CoRR, abs/2105.00 00000000000 2021.
[26] Xiaorui Liu, Wei Jin, Yao Ma, Yaxin Li, Hua Liu, Yiqi Wang, Ming Yan, and Jiliang Tang. Elastic graph neural networks. In International Conference on Machine Learning, ICML, 2021.
[27] Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang. SOFT: softmax-fre transformer with linear complexity. In Annual Conference on Neural Information Processing Systems, NeurlPS, 2021.
[28] Songtao Lu, Mingyi Hong, and Zhengdao Wang. PA-GD: on the convergence of perturbed alter- nating gradient descent to second-order stationary points for structured nonconvex optimization. In International Conference on Machine Learning, ICMI., 2019.
[29] Yao Ma, Xiaorui Liu, Tong Zhao, Yozen Liu, Jiliang Tang, and Neil Shah. A unified view on graph neural networks as graph signal denoising. CoRR, abs/2010.01777. 2020.
[30] Andrew Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Annual Meeting of the Association for Computational Linguistics, ACL, 2011.
[31] S. Deepak Narayanan, Aditya Sinha, Prateek Jain, Purushottam Kar, and Sundararajan Sella- manickam. IGLU: efficient GCN training via lazy updates. CoRR, 0ttp21099.1999 2021.
[32] Xuran Pan, Shiji Song, and Gao Huang. A unified framework for convolution-based graph neural networks, 2021.
[33] Jeffrey Pennington, Richard Socher, and Christopher D. Manning. GloVe: Global vectors for word representation. In Conference on Empirical Methods in Natural Language Processing, EMNLP, 2014.
[34] Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosFormer: Rethinking softmax in attention. In The Tenth International Conference on Learning Representations, ICLR, 2022.
[35] Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitner, Thomas Adler, David P. Kreil, Michael K. Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. In International Conference Learning Representations, ICLR, 2021.
[36] R. Tyrelll Rockafellar. Convex Analysis. Princeton University Press, 1970.
[37] Pierre Samuel. Projective Geometry. Springer, 1988.
[38] Mark Schmidt, Nicolas Le Roux, and Francis R. Bach. Convergence rates of inexact proximal- gradient methods for convex optimization. Annual Conference on Neural Information Processing Systems, NeurIPS, 2011.
[39] Han Shi, Jiahui Gao, Hang Xu, Xiaodan Liang, Zhenguo Li, Lingpeng Kong, Stephen Lee, and James Kwok. Revisiting over-smoothing in BERT from the perspective of graph. International Conference Learning Representations, ICLR, 2022.
12

---

[40] Richard Socher, John Bauer, Christopher D. Manning, and Andrew Y. Ng. Parsing with compositional vector grammars. In Annual Meeting of the Association for Computational Linguistics,, ACI., 2013.
[41] Pablo Sprechmann, Alex Bronstein, and Guillermo Sapiro. Learning efficient sparse and low rank models. IEEE Trans. Puttern Analysis and Machine Intelligence, 37(9), 2015.
[42] Ying Sun, Prabhu Babu, and Daniel Palomar. pooorizationmminmmzzan algorithms in signal processing, communications, and machine learning. IEEE Transactions Signal Processing, 65(3):794-816, 2016.
[43] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Annual Conference Neural Information Processing Systems, NeurlPS, 2017.
[44] Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference Learning Representations, ICLR, 2018.
[45] Zhangyang Wang, Qing Ling, and Thomas Huang. Learning deep to encoders. In AAAI Conference on Artificial Intelligence, 2016.
[46] Lin Xiao and Tong Zhang. proximal stochastic gradient method with progressive variance reduction. SIAM J. Optim., 24(4):2057-2075. 2014.
[47] Xingyu Xie, Qiuhao Wang, Zenan Ling, Xia Li, Yisen Wang. Guangcan Liu, and Zhouchen Lin. Optimization induced equilibrium networks. CoRR, abs/2105. 13228, 2021.
[48] Chaoqi Yang, Ruijie Wang, Shuochao Yao, Shengzhong Liu, and Tarek F. Abdelzaher. Revisiting "over-smoothing" in deep GCNs. CoRR, abs/2003. 13663, 2020.
[49] Yongyi Yang. Tang Liu, Yangkun Wang, Jinjing Zhou, Quan Gan, Zhewei Wei, Zheng Zhang, Zengfeng Huang, and David Wipf. Graph neural networks inspired by classical iterative algorithms. In International Conference Machine Learning, ICML, 2021.
[50] Yongyi Yang, Yangkun Wang, Zengfeng Huang, and David Wipf. Implicit unfolded graph neural networks. CoRR, abs/2111.06592, 2021.
[51] Alan L. Yuille and Anand Rangarajan. The concave-convex procedure (CCCP). Annual Conference Neural Information Processing Systems, NeurlPS, 2001.
[52] Hongwei Zhang, Tijin Yan, Zenjun Xie, Yuanqing Xia, and Yuan Zhang. Revisiting graph convolutional network on semi-supervised node classification from an optimization perspective. CoRR, abs/2009. 14699, 2020.
[53] Kaixiong Zhou, Xiao Huang, Daochen Zha, Rui Chen, Li Li, Soo-Hyun Choi, and Xia Hu. Dirichlet energy constrained learning for deep graph neural networks. CoRR, abs/2107.02392 003992, 2021.
[54] Meiqi Zhu, Xiao Wang, Chuan Shi, Houye and Peng Cui. Interpreting and unifying graph neural networks with an optimization framework. CoRR, abs/2101.11859, 2021.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A] Our work is mostly theoretical and devoted to the better understanding of an existing, well-
established modeling framework; we do not see any clear pathway to negative societal impact.
13

---

(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] They in the supplementary.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g.. with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [N/A]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you deseribe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "yEqVIuUU7-h", "reviewer": "Reviewer_9kiY", "summary": "This paper proposes to understand Transformer through the lens of unfolding optimization. It overcomes several challenges in previous works of formulating the complex Transformer architecture and obtains several insights from the new interpretation.", "questions": "Please refer to the weakness part.", "limitations": "The author may want to elaborate on the implication of the proposed interpretation.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The formulation of the Transformer layer as unfolding optimization is novel in the Transformer literature and brings several new ideas, e.g., potential improvement of the architecture inspired by optimization techniques.\n- The paper clearly states and tackles the challenge from previous works and is able to formulate the whole complicated Transformer layer.\n\nWeakness:\n- It would be great if the authors can further elaborate on how can we benefit from this new interpretation, e.g., any direct implication on architecture design.\n- Is unclear if this analysis can be extended to a more practical setting of Transformer models, i.e., with components like LayerNorm.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "j9CYAlnrTB", "reviewer": "Reviewer_siVJ", "summary": "To find an energy function underlying the Transformer model, the authors first outline 4 major obstacles and focus on solving or at least partly partially addressing them. The experimental results also help to support their results.", "questions": "How does the energy function underlying the Transformer help us to design or train the transformer? I’m interested in how useful the theorems are.", "limitations": "NA", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\nThis paper is well-written. The background is easy to understand and friendly to the readers not in the unfolded optimization track. The problems the authors want to solve are also well-defined and explained clearly.\nThe novelty is also well claimed. The authors provided a clear comparison between this work and existing literature.\nThe experiments are also helpful to verify the viewpoints in this work.\n\nWeakness:\nThe transformer architecture has been simplified. I worry about the impact of normalization layers and residual connection. Also, the order of different layers may also have an impact to the results. For instance, what would happen if the transformer used pre-LayerNorm or post-LayerNorm?\n\nSimilarly, what would happen if we use the widely used multi-head attention and other non-linear activation functions (e.g. GeLU)? The corresponding analysis is missing in this work.\n\nIn real-world applications, we usually use a pre-trained transformer instead of Glove+Transformer. Authors’ evaluation setting is a bit different from the real-world ones.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "FIvF6W64rK1", "reviewer": "Reviewer_zDqk", "summary": "This paper tried to find an energy function underlying the transformer model and reinterpret transformers as the unfoling of an interpretable optimization process across iterations. This paper starts from the complete Transformer architecture, and due to the structural complexity of the transformer, this paper addresses several challenges in the derivation process including heterogeneous layer types and non-linear activations.", "questions": "Please refer to the weaknesses.", "limitations": "No but not need.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- Detailed theoretical proofs are provided.\n- Writing is logical and well laid out.\n\nWeaknesses\n- The approach of this paper seems to construct an energy function behind the network optimization process by analyzing the existing standard Transformer structure. Although the article mentions the practical value, this approach seems to be difficult to inspire the network design and the experimental part fails to provide the corresponding proof. Is there any actual demonstration of constructing a network out of the energy function in reverse?\n- This paper mentions that the unfolding perspective has been adopted to analyze MLPs and CNNs. In terms of comparison at the network level, is there an opportunity to compare the energy function behind the different networks and the difference between their strengths and weaknesses?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "D3d2XkAZU6S", "reviewer": "Reviewer_4boQ", "summary": "The paper attempts to analyze the Transformer architecture from energy minimization perspective. In other words, authors seek an energy function (of input representations), such that minimizing this function is equivalent to running the forward pass of a corresponding transformer model. The paper lists the key challenges of building such energy function. Authors then address these challenges and propose an AIM-based procedure that corresponds to running the forward pass of a transformer with some additional constraints on model parameters.\nFinally, the paper conducts two practical experiments where they show that transformers \"as described previously, up to known limitations \nlike symmetric weights\" actually reduce the corresponding energy function with subsequent model layers.\n\n", "questions": "\n### Minor corrections:\n> (supplementary) ./code/README.md L5: python appllo_circle.py\n\nperhaps, `apollo_circle` ? (typo)\n\n\nThough authors have every right to decide for themselves, I would also recommend specifying the environment needed to run the supplementary code. Once could, for instance, add requirements.txt with specific library versions, or a docker container. This would simplify the task of reproducing the paper's results, especially in the far future.\n\n", "limitations": "The paper explicitly states the main limitations of their theoretical results (i.e. the requirement for symmetric weights).\nHowever, I believe that the limitations of the proposed empirical results could be stated more clearly, so as not to mislead practically minded readers. I elaborate on these limitations in the \"Strengths And Weaknesses\" section.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Authors propose a curious and novel (to my understanding) perspective on the popular transformer architecture.\nTo the best of my understanding, the overall approach looks sound, and the specific theoretical derivations for constructing the optimization procedure look solid. However, I am not an expert in the narrow domain of viewing and may have missed something.\n\nMy main concerns, however, are about practical applicability of the derived results.\n\n### Section 6, Figures 5,6\n\nBased on my understanding of the supplementary code (in particular, ./code/energy_curve/main.py , line 32-55) and the text presented in Section 6, __the energy figures 5,6 represent an untrained model.__\n\nFor instance, the figure with relu is constructed by the following code in ./code/energy_curve/main.py :\n\n```\nmodel_2 = Transformer(d , num_layers , 2)\nmodel_2.normalize_weight()\npaint(model_2 , \"../generated_figures/%s.png\" % relu_name)\n```\n\nThe model is initialized, normalized, and immediately used for producing the plot. The training data is loaded a few lines above, but ultimately does not affect the plot in any manner except one: the model input dimension is based on word embedding (d). However, the embeddings themselves do not participate in updating model parameters.\n\nInterestingly, __Section 6 does not explicitly contradict that__, but does not discuss it either. As such, the reader may be misled into believing that the practical experiments are more significant than they actually are.\n\nTo the best of my understanding, the actual evidence of practical applicability of this work is as follows:\n\n1. all observations are made not with actual transformers, but with the proxy as defined in the paper (stated in Section 6)\n    - in addition to the assumptions stated in the paper, there are other subtle differences, i.e. lack of scaling, or unconventional initialization\n2. all observations are mad with randomly initialized weights, without any training on the data\n3. it is unclear whether the said proxy retains the practical properties of actual transformers, i.e. whether it can actually solve the same tasks with the same quality\n4. if it does, it is unclear whether the properties of the proxy transformer reflect the properties of actual transformer __once it has been trained on actual data__\n\nWith these four concerns in mind, it might be premature to claim that:\n\n> Hence although seemingly complicated, the conditions adopted by Lemma 4.3 and Theorem 5.1 are nonetheless likely hold in many practical settings\n\n> the unfolding perspective can be used to devise architectures with inductive biases aligned with downstream tasks to help compensate for less model flexibility\n\nTo summarize, I believe this paper proposes a curious theoretical perspective that may yet improve our understanding of real-world transformers. However, in its current for, it is unclear how well does the proposed approximation reflect the real world behavior of transformers, especially when trained.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept