Convexifying Transformers: Improving optimization and understanding of transformer networks
Authors: Tolga Ergen
Behnam Neyshabur
Harsh Mehta
Venue: Submitted to ICLR 2023 (2023)
Keywords: Convex optimization
transformers
attention
self-attention
group sparsity

Abstract: Understanding the fundamental mechanism behind the success of transformer networks is still an open problem in the deep learning literature. Although their remarkable performance has been mostly attributed to the self-attention mechanism, the literature still lacks a solid analysis of these networks and interpretation of the functions learned by them. To this end, we study the training problem of attention/transformer networks and introduce a novel convex analytic approach to improve the understanding and optimization of these networks. Particularly, we first introduce a convex alternative to the self-attention mechanism and reformulate the regularized training problem of attention/transformer networks. Then, we cast the reformulation as a convex optimization problem that is interpretable and easier to optimize. Moreover, as a byproduct of our convex analysis, we reveal an implicit regularization mechanism, which promotes sparsity across tokens. Therefore, we not only improve the optimization of attention/transformer networks but also provide a solid theoretical understanding of the functions learned by them. We also demonstrate the effectiveness of our theory through several numerical experiments.

Under review as a conference paper at ICLR 2023
CONVEXIFYING TRANSFORMERS: IMPROVING OPTI- MIZATION AND UNDERSTANDING OF TRANSFORMER NETWORKS
Anonymous authors Paper under double-blind review
ABSTRACT
Understanding the fundamental mechanism behind the success of transformer net- works is still an open problem in the deep learning literature. Although their remarkable performance has been mostly attributed to the self-attention mecha- nism, the literature still lacks solid analysis of these networks and interpretation the functions learned by them. To this end, we study the training problem attention/transformer networks and introduce novel convex analytic approach improve the understanding and optimization of these networks. Particularly, we first introduce convex alternative to the self-attention mechanism and reformu- late the regularized training problem of transformer networks with our alternative convex attention. Then, lem that is interpretable and easier to optimize. Moreover, as byproduct of our convex analysis, we reveal an implicit regularization mechanism, which promotes sparsity across tokens. Therefore, we not only improve the optimization of atten- tion/transformer networks but also provide a solid theoretical understanding of the functions learned by them. We also demonstrate the effectiveness of our theory through several numerical experiments.
the reformulation
convex optimization prob-
INTRODUCTION
Transformer networks proposed by Vaswani et al. (2017) have become a dominant architecture in various tasks, especially Natural Language Processing (NIP) (Devlin et al., 2018; Radford et al., 2019), due to their extraordinary generalization properties and high capacity to learn from vast amount data. Although there exists substantial empirical evidence on the effectiveness of trans- former networks, revealing the underlying theoretical reasons behind their success is still open research problem due to their highly nonlinear and nonconvex structure.
A significant body of research focused on analyzing certain components of transformer networks via empirical studies. As an example, Liu et al. (2021a); Vashishth et al. (2019); Dong et al. (2021); Voita et al. (2019); Takase et al. (2022); Liu et al. (2021a) studied the impact of the attention mech- anism on transformer networks. Although these studies agreed that attention an essential com- ponent of transformers, they also raised several issues regarding interpretability and optimization. Particularly, Voita et al. (2019) demonstrated that most attention heads can be removed without af- fecting the performance of the network, which is an indicator of large amount of redundancy in the network. Vashishth et al. (2019) provided set of empirical evidence showing that attention might be needed for some NLP tasks. Additionally, Dong al. (2021) revealed that although atten- tion is at the heart of transformer networks, training attention network in the absence of Fully Connected Network (FCN) layers and skip connections is extremely challenging since the network output degenerates quickly without them. Similarly, Takase et al. (2022) discussed the importance of layer normalization and skip connections for transformer networks so that even changing the po- sition of these might considerably impact the performance of transformer network. However, solid theoretical analysis of the underlying factors behind these issues is sill lacking, likely due to the highly complex and nonconvex structure of transformer networks.
A series of papers also focused on designing new alternatives to the self-attention mechanism which perform similarly and might provide further interpretations towards the overall model. One set of

---

Under review as a conference paper at ICLR 2023
Our Convexified Transformer Block
Simplified Transformer with proposed attention
Standard Transformer
FCN
FCN
Activation
X
Regularization: Sparsity across takens
Add Norm
+
Multi-head Attention
Convex Combination
Linear Layer
Tokens
Tokens
Tokens
Figure I: Summary of our main findings: We first propose an alternative attention, i..., taking the convex combinations of tokens, and then convexifying whole transformer block (attention + Fully Connected Network (FCN) with this new attention mechanism. The equivalent convex formulation also reveals a sparsity-iduucing regularization across tokens as detailed in Theorem 2, and 3.
work utilizes multi-layer perceptron based architectures, Tolstikhin et al. (2021); Tatsunami & Taki (2021); Touvron al. (2021); Liu al. (2021b); Yu et al. (2021), while another set of of papers pro- poses Fourier based models Lee-Thorp et al. (2021); Rao et al. (2021); Li et al. (2020); Guibas et al. (2021). Others also proposed replacing the self-attention mechanism with matrix decomposition Geng et al. (2021). Although these works successfully applied to certain applications, they lack any solid theoretical analysis and understanding from an optimization perspective. Recently, Sahiner et al. (2022) attempted to analyzed transformer networks via convex duality by completely chang- ing structure the self-attention mechanism and removing FC layers. Even then, they failed to provide solid practical implicationss/enefitss for transformers since their formulations are extremely challenging and complex to be solved in practice.
Recently, another line of research has focused understanding structures and patterns emerge throughout the training of transformer networks (Power al., 2022; Thilak et al., 2022; Barak al., 2022). In particular, the grokking phenomenon was first observed by Power al. (2022) on specific algorithmic tasks, such as modular division operations. Specifically, grokking refers to sudden transition of validation or test accuracy perfect generalization and this generalization happens well past the point of perfect training accuracy. This interesting behavior contradicts the common practice of early stopping in the training of deep learning models and definitely requires further understanding as to why this phenomenon emerges.
In order to remedy the issues associated with the standard transformer networks, in this paper, we develop a convex optimization perspective to train, analyze and understand transformer networks. Particularly, we first propose a convex alternative the self-attention mechanism and then develop our convex analytic framework on the resulting model as detailed in Figure 1.
1.1 CONTRIBUTIONS
Our contributions can be summarized as follows:
We propose an alternative formulation to the standard self-attention mechanism and study the regularized training problem of attention/rransfommer networks with
We convexify the regularized training problem of attention/transfommed networks with the pro- posed attention layer as shown Figure and therefore enable finding a globally optimal solu- tion without requiring any nonconvex optimization heuristic, e.g., layer normalization and skip connections.
We also apply convex analytic framework to various architectures, e.g., networks with or without an FCN layer. Thus, we are able to explain the impact of each component on the models learned throughout training.
We reveal an implicit regularization mechanism induced by our attention mechanism. We further characterize this regularization sparsity-inducing factor across tokens.
2

---

Under review as a conference paper at ICLR 2023
We demonstrate the effectiveness of our reformulation via various experimental results. We also show that our reformulation sigmificantly mitigates the grokking phenomenon studied in recent papers (Power et al., 2022: Thilak al., 2022).
1.2 NOTATIONS
We use lowercase and uppercase bold letters to denote vectors and matrices, respectively. We denote certain column/element of a vector or matrix using subscripts. As an example, Wjj denotes the jk entry of the matrix W. We use I& to denote the identity matrix of size k k and 0 to denote vector/matrix zeros (or ones) with ap- propriate sizes. We also use [n] for the set of integers rang- h
Table 1: Notations.
Notation
Description
N
# of sentences/samples # of tokens embedding dimension
# of heads
ing from to n. We represent the Fuclidean and Frobe- nius norms as || ll2 and ||r, respectively. We also use [x 2 0J to denote the valued indicator function. We provide more notations we use throughout the paper in Table 1.
# of outputs
TRANSFORMER NETWORKS
Given data sample (or sentence) X E Rexd as sequence of h tokens with the embedding dimen- sion d, we define the key, query, and value matrices as
Q XW
Wg Rdxd
K XW
k:
W E Rdxd
V = XW.,
W. E Rdxd
which are the main components of the self-attention mechanism. Then, single transformer block, which is basically a stack self attention, residual connection, layer normalization, and point-wise feedforward connections, can be formulated as follows
A.. softmax (QK') V
A, A, W.. W. E Rddd
(I)
X LayerNorm (A,) +X
X WX, We
where denotes the activation function for the FCN layer. Although skip connections, layer normalization and FCN also play crucial role in a transformer block, the success of these networks has been mostly attributed to the self-attention part, denoted as A, (Vaswani al., 2017). Therefore, in the following section, we first study the training problem of simplified transformer network, for which the network output directly A.. We then extend our derivations transformer network with FCN layers.
ATTENTION-ONLY NETWORKS
We first consider simplified transformer network only with self attention layer that maps input sequence X E nxd to the output sequence Ý Ra with outputs as follows
Ý = softmax ( W x XW.W
We also call the model (2) attention-only network. This is meaningful model and has been applied to various tasks, including machine translation, language modeling, image captioning. and object recognition (Vashishth et al., 2019).
We next consider standard regression framework with an arbitrary convex loss function. Given a training set (X,, Y,) where X; E Rexd and Y, E Raax denote the input sequence and the

---

Under review as a conference paper at ICLR 2023
labels/target outputs, respectively, the weight decay regularized training problem for the attention- only network (2) is as follows
N
3
Σ c (softmax (XW, XW. Y,)
Σ
I # || (3)
min W.WWE.W..WW
#E(g.k.
where C (-) is an arbitrary convex loss function, including squared loss and cross entropy, and 3 is the regularization coefficient.
Although the attention-only model in (2) is quite powerful across various NLP tasks, e.g., natural language inference, neural machine translation, and text classification (Vashishth et al., 2019), the corresponding training problem in (3) is an extremely challenging optimization task and requires various nonconvex optimization heuristics (Dong et al., 2021) to be adequately trained. To rem- edy these issues, in the following sections, the attention part with an alternative convex layer and then cast the reformulated training problem interpretable convex optimization problem that enables the globally optimizing the network parameters.
first reformulate the training problem by replacing
3.1 CONVEX ATTENTION LAYER
We first note that since the softmax (+) operation highly nonlinear and nonconvex, the training problem in (3) is a challenging nonconvex optimization problem. Therefore, one may not adequately train attention networks and obtain trivial models at the end of training. For example, Dong et al. (2021) shows that attention networks are likely to degenerate throughout the training and the output converges to a rank-1 matrix. Thus, they fail to learn the underlying tasks.
To avoid the issues associated with the nonconvex formulation in (2), first replace the softmax operation with a simpler yet effective alternative. Particularly, since softmax converts the rows of its input matrix to probability distribution, can be relaxed linear operation with unit simplex constraints as follows
for any U E RuXn EW E s.t. softmax (U) XXW,,
1, vi [n|} denotes convex of constraints, also termed as unit simplex constraints. Thus, we simplified and convexified the attention mechanism without disturbing structure. Based on this observation, (3) can be reformulated as follows
where A (W Raxa w; 2 0, 0,1 w;
3 (1W2l + |aal
min W,EA
<wW Y.)
(4)
W. W 1=1
Note that the model above utilizes single head attention model and, therefore, may not be practi- cally relevant due to its insufficient expressive power. Thus, we introduce the concept of head to the problem in (4) as follows
N
h
min WIJEA W Exdx W ERdx
Σ /=1
w
W, W3y. Y, + Σ W2 + Warl
1=1
(S)
Now, we are ready apply the convex analytic tools (5) as detailed in the next section.
3.2 CONVEX OPTIMIZATION FOR ATTENTION-ONE NETWORKS
As a warm-up, let us consider the scalar output prediction problem where the targets are one- dimensional, i.c., Ii E R. Then, (5) reduces to the following optimization problem
h
min WIFA
Σe c
Σ
w X.W. Vi
+
(1 (way)?).
(6)
2
W2 ER
Next, we first apply rescaling between the parameters Waj and Way such that (6) can be described & regularized optimization problem.
4

---

Under review as a conference paper at ICLR 2023
Table Number parameters and FLOPs for the convex and nonconvex models. Here, we use the following notations: n. # of tokens, d: embedding dimension, h: # of heads, and c: # of outputs.
Nonconvex
Convex
Standard
Alternative (Ours)
# params
FLOPs
of params
FLOPs of params
FLOPS
Scalar output Multi output Multi output with FCN
h33d"
+l+++111
Oind]
nd
O(ud)
A33d + dej
O(n'dh nahe)
Onna
nde ndeb
O(dde ((ndeb)
M33 de)
O(n²dh ndhe)
M(n+ d++)
O(nd+t))
Lemma 1. The problem in (6) is equivalent to the following & regularized training problem
N
min W
Σ
W22 + 3|walll
Way ll Slargy ER
j=1
Based the equivalent formulation in Lemma 1, the next theorem introduces a convex optimization problem that is equivalent to (6).
Theorem 1. The nonconvex optimization problem (6) can be equivalently cast the following convex optimization problem
min ZERoXd
Σ c (trace
Ui) + Σ 1zxll2
Note that the equivalent convex model (8) requires a single parameter matrix Z E Rexd where each row is the attentions scores of the corresponding token. We also remark that the regularization in (8), i.e., the sum of & norms of the rows of the parameter matrix Z, a specific type of regular- ization, also known group & Lasso, introduced by (Bakin al., 1999) and shown promote group sparsity across parameters (Yuan & Lin, 2006). In our case, the group sparsity is across the token index k. Therefore. one can interpret the model (8) as sparse linear model, where the sparsity is across tokens. In other words, (8) can be explained as a model that tries to use as few tokens as possible to fit the training labels {y: {
Unlike the nonnegative attention scores in (6), denoted as Wiy E A, the convex parameters Z e Raxd do not require any constraints. Therefore, one can directly apply standard training algorithms, such as SGD and Adam, to train the convex problem (8). Moreover, for (6) can be recovered from a solution to (8) as proven in the following result.
optimal set of parameters
Proposition 1. After solving the convex optimization problem in (8), one can recover an optimal solution to the noncorvex optimization problem (6), denoted as {wig. Why wG, > as follows
Viz, 12 Wa. V1lzal22: [h],
wi, C Waj
where e, E 1.th is the 1th ordinary basis vector, zj E Rd is the jth row of Z, and we assume that there are nonzero rows out of rows of due to the aparsity-inducing regularization (8).
Proposition 1 proves that there a one-to-one mapping between the parameters of the nonconvex formulation in (6) and the convex formulation in (8). Therefore, there is need solve the chal- lenging nonconvex optimization problem (6) which also requires several optimization heuristics to be adequately trained. Instead, one can solve the convex problem (8) and then use the mapping in Proposition to obtain an optimal solution to (6).
3.3 EXTENSION TO MULTIDIMENSIONAL OUTPUTS
In the previous section, we considered setting with scalar target variables, i.c., Yi R. How- ever, for some problems, e.g., multiclass classification, target variables can be multidimensional. Therefore, we now extend the analysis to the problems with multiple/vector outputs as follows
8 [[ 112 + || I1)
min
C
wi X,
X, W3j.
(9)
WIJE WayER ERy
2
1=1
5

---

Under review as a conference paper at ICLR 2023
where y; E R and denotes the number of outputs/classes. Note that here we put E fi-norm Waj to enable our convex arguments but this does not impact performance of the network in practice. Then, following the same derivations yields the convex program in the next result.
Theorem 2. The nonconvex optimization problem (9) is equivalent to the following convex opti- mization problem
min ZIERaXd
c (trace (.) ya) +
1znll2
(10)
Theorem shows that the equivalent convex model becomes separable over the output index 1, i.e., instead of single parameter matrix in (8), here we have parameter matrices due to having coupputs the nonconvex model (9) (see Table for details). This also illustrates that the number of outputs in the network directly controls the overparameterization level of the equivalent convex formulation.
3.4 ATTENTION NETWORKS WITH FCN LAYERS
Although the model in (5) exhibits interesting properties in various applications (Dong et al., 2021), I), is basically linear function of the token matrix X. Therefore, it likely suffer from inadequate performance especially for some challenging problems in NLP. A series of papers (Dong et al., 2021; Geva et al., 2021; Meng et al., 2022; Geva et al., 2022b;a) also confirmed the importance of FCNs via extensive empirical evidence. Therefore, in this section, we include an FCN layer to our attention only model in (5) and derive an equivalent convex formulation for this new model.
Here, we consider the following optimization problem
N
h
min EA
C
Σ
w 1i
X. W22
Waj.
+
1 ||
(II)
2
Way.Way ER
where a (-) is the activation function.
Theorem 3. The nonconvex optimization problem (11) with the gated ReLU activation is equivalent the following convex optimization problem
N
h Σ
h
min ZRgRax
1sytaaee (Z..X.) ya
+
12
(12)
1=1
where 1; = (uT, 2 0} denotes the indicator function for gated ReLU activation and here
{uij, 1 are fixed vectors that can be randomly selected.
Theorem 3 implies that introducing the activation function further increases in the overparameteri- zation level of the equivalent convex formulation. Precisely, (12) has h times more parameters than (10) as shown in Table 2.
4 NUMERICAL EXPERIMENTS
In this section, we present experimental results corroborating our theory in the previous sections.
Student-feacher setting with BERT:
We first consider a student-teacher setting with the pretrained Particularly, we feed
BERT model in the Hugging Face repository, i.e., bert the samples from the mrpc subset of the glue dataset (Warstadt al., 2018; Wang al., 2019) through the pretreained BERT model and save the input and output activations in certain layer. Then, we train the attention-only models, i.c., standard nonconvex self-attention (3), alternative nonconvex attention (9), and convex (10), from scratch using these pre and post activations as training dataset. All the experiments throughout this section are performed using single GPU on Google Colab. We also use the same regularization coefficient 8 and optimizer, i.c., Adam, and tune the learning rate and regularization coefficient by performing a grid search on a validation dataset for both algorithms. However, notice that we do any nonconvex optimization heuristics, e.g..

---

Under review as a conference paper at ICLR 2023
Objective value
10-1
Convex(Ours)
Test error
Nonconvex Stannaad
10-1
Noncorvex Altemative
10
10-1
100
150
200
300
0
100
200
250
300
Time (s)
Time (s)
(a) Objective value
(b) Test error
Figure 2: Comparison of the convex and nonconvex models on the dataset extracted from pretrained BERT architecture in student-teacher setting. Here, include two non-conver models. specifically standard self-attention model KNonconvex-Standan in (3) and alternative attention (Nonconvex- Alternative) in (9). Our convex training approach achieves significantly lower objective value and test error than the original nonconvex training.
0.7
0.000
0.004
0,055
20
20
20-
00016
0.3
40
0.2
40-
6.005
0.000
10
20
30
4.0
50
10
20
30
40
50
10
20
30
50
(a) Ground truth
(h) Nonconvex
(c) Convex
Figure 3: Attention maps obtained by our convex and standard nonconvex training approaches as well the ground truth attention map in the BERT model. Here, show that the nonconvex training fails to learn the underlying patterns and simply achieves uniform attention map. However, our convex training approach outputs attention map that is close the ground truth map.
layer normalization and skip connections, for the convex model all the experiments. Figure 2, we plot the objective values (i.e. training loss regularization term) and test losses with respect to time in seconds using the data extracted from the sixth layer of pretrained BERT model. We observe that our convex training approach achieves almost order of magnitude smaller objective value than standard nonconvex training, which possibly stuck at local minimum. This effectiveness in training also translates into better generalization, i.c., convex training approach obtains lower test loss than the standard nonconvex training. order to understand the functions learned by each models, also analyzed the attention maps in Figure 3. Here, standard nonconvex training fails learn the underlying model and outputs uniform attention map across token. However, our convex training outputs an attention map that quite similar the ground truth attention map, and therefore we successfully learn the structure in the training data. Hence, these experiments clearly illustrate the effectiveness of our convex training approach in both training and testing.
A Algorithmic datasets and Grokking: Inspired by the grokking phenomenon observed in Power al. (2022), we next validate the effectiveness of our convex training approach against standard transformer networks with the self-attention mechanism in (1) on algorithmic datasets. Particularly, we use the same setting in Power et al. (2022), and evaluate the performance on modular division operations with mod 97 and mod 15, where we train the architectures till they reach 99% test accuracy whenever possible. In Figure 4, we first replicate the results in Power et al. (2022) and confirm that the grokking phenomenon indeed emerges here, i.e., the nonconvex curve (purple) reaches 100% training accuracy at around 10 iterations in Figure 4a while it requires more than 10S iterations reach perfect generalization in Figure 4b. We also compare the nonconvex and convex training approaches. Here, we show that our convex training approach converges to the perfect generalization accuracy 10x faster than the nonconvex one in Figure 4b. Moreover, the

---

Under review as a conference paper at ICLR 2023
1.0
1.0
Accuracy
Accuracy
0.0
Loss
10:1
0.4
0.4
10:2
0.2
10-3
RenCHTWe
10:1
10a
107
10
10
10
10
10'
301
10r
10a
10
10"
13t
10
Iterations
Iterations
Iterations
(a) Training accuracy
(b) Test accuracy
(c) Test loss
Figure 4: Comparison of the convex and nonconvex models on the modular division operation mod 97. Here, we train the networks to reach 99% test accuracy and show that our convex training approach exhibits significantly faster convergence and lower test loss than the nonconvex training.
1.0
Noncenvex (L=1)
10a
Accuracy
Nonconvex
10a
Convex (L=1)
0.6
Convex (L-2)
Loss
10
0.4
10
0.2
10-3
10-1
0.0
10a
10?
10
10
10
101
102
107
10
10*
Iterations
Iterations
(a) Test accuracy
(b) Test loss
Figure 5: Comparison of one- and two-layer transformer networks on the modular division task mod 97, where L denotes the number of layers in each model. We observe that introducing one more layer substantially improves the convergence speed of our convex formulation while it fails to make noticeable impact the nonconvex formulation.
convex model also yield significantly lower test loss in Figure 4c, which implies that it has higher confidence test predictions and therefore more robust than standard nonconvex training.
We remark that in the previous section, we theoretically analyze only single attention/rransfommer blocks. However, since the benign impact of depth or number of layers (denoted as L) has already been empirically proven in the deep learning literature, we also propose an extension of our convex model deeper settings. We basically stack the convex transformer layers in (12) to obtain an arbi- trarily deep network. In Figure 5, we compare the performance of two-layer transformer networks with one-layer networks. Here, we observe that while adding one more layer results in significant improvements for the convex model, especially in terms of optimization speed, fails to make any discernible difference for the nonconvex model. Moreover, we run the algorithms on the mod 15 operation which is basically more challenging task due smaller number of samples. In this case, one-layer models are not able learn the underlying task perfectly as demonstrated in Figure but our convex model is significantly better in terms of both test accuracy and test loss. By increasing the number of layers to four, we enable both models to achieve perfect generalization accuracy. Our deep model reaches this level much faster and also yields lower test loss than the nonconvex model.
We next empirically analyze the grokking phenomenon on both our convex and standard nonconvex models. For this purpose, we plot the number of iterations to reach 99% test accuracy for each of our experiments in Figure 7a. Notice that here we do not include the one-layer results for the mod 15 case, since both models fail to achieve perfect generalization in that case. Figure 7a clearly shows that convex training approach converges to the 99% accuracy level substantially faster than the standard nonconvex training. Therefore, we also mitigate the impact of the grokking phenomenon demonstrated in Figure 7b, where quantify the amount of grokking in terms of the number of iterations. Based on this experiment, we also conjecture that the grokking phenomenon can be mostly attributed to the highly nonlinear and nonconvex structure of standard transformer models.
8

---

Under review as a conference paper at ICLR 2023
Noncorivex (L=1)
10
Accuracy
0.8
Nonconvex (Lm4)
Convex (L=1)
0.6
Convex (L-4)
Loss
10-1
0.2
10
0.0
10-3
10
10
10
10
10a
10a
10a
10
10r
10
Iterations
Iterations
Training accuracy
(b) Test loss
Figure 6: Comparison of one- and four-layer transformer networks the modular division task mod 15. Here, one-layer networks fail to achieve 99% test accuracy however our convex training approach (light green) still generalizes better than the nonconvex training (light purple). We also show that perfect generalization terms of accuracy can be achieved with four-layer networks.
# of iters
10
Grokking
10a
10a
10a
(15.4)
(97.1)
(97.2)
(15.4)
(97.1)
(97.2)
(p,L)
Convex(Ours)
(p,L)
Nonconvex
(a) Training iterations
(b) Grokking iterations
Figure 7: Amount grokking in terms of # of iterations required by our convex and standard nonconvex training approaches, where p and denote the coefficient for the modular division and the number of layers, respectively. Here, we do not include the one-layer results in Figure 6 since both algorithms fail achieve 99% test accuracy. We demonstrate that the impact of grokking is substantially mitigated with our convex training approach.
CONCLUSION
In this paper, we studied the regularized training problem of attention/rransformmr networks and developed convex analytic framework to train these networks. Particularly, first proposed convex alternative to the self-attention mechanism and then reformulated the training problem with this alternative attention mechanism convex optimization problems. Thanks to mulation, we globally optimize the network parameters without requiring any kind of nonconvex optimization heuristics. In addition, the functions learned by reformulation is transparent and interpretable. More importantly, the reformulated problem reveals a sparsity-inducing -induiing regularization mechanism across tokens the data, which also sheds more light on the structure of the resulting function and its generalization properties. We then empirically verified effectiveness of our convex training approach over standard nonconvex training via several numerical experiments.
refor-
We also note that analyzing transformer networks through the lens of convex optimization theory is extremely crucial since may result substantial improvements in the understanding and optimiza- tion of these networks. However, is also quite challenging due to the inherent nonconvex structure of the network model. To the best of our knowledge, this paper is the first step in this direction and therefore has some limitations which can hopefully be eliminated by future work. Specifically, in this paper, we mainly focused on the theory side of convex analysis and empirically validated the theory few small-scale problem instances. We hope that comprehensive and large-scale empirical verification of our theory will be conducted by the follow-up papers.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):6296681) 2017.
Sergey Bakin et al. Adaptive regression and model selection in data mining problems. 1999.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. arXiv preprini arXiv:2207.08799, 2022.
Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Jacob Devlin, Ming- Hiig Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: /800vv//00004000 2018.
Yihe Dong, Jean- -aa-itat Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth, 2021. URL https: http://axxvvv..o/ abs 2003/2103.0 03 404.
Tolga Ergen and Mert Pilanci. Convex geometry and duality of over-parameterized neural networks. Journal of machine learning research, 2021.
Zhengyang Geng, Meng-Hao Guo, Hongxu Chen, Xia Li, Ke Wei, and Zhouchen Lin. Is attention better than matrix decomposition? arXiv preprint (tti::2209.045550 2021.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen- tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021, pp. 5484-5495, Association for Computational Linguistics, 2021. doi: 10.18653/v1/2021. emnlp-main.444 446. URL https: / /ooi .rgg 10. /v15553v1/222
emn p wewnl-mmaee i 446 6.
Mor Geva, Avi Caciularu, Guy Dar, Paul Roit, Shoval Sadde, Micah Shlain, Bar Tamir, and Yoav Goldberg. Lm-debugger: An interactive tool for inspection and intervention in transformer- language models. arXiv preprint arXiv: Arxiv-2204. /2/30, 2022a.
Mor Geva, Avi Caciularu, Kevin Ro Wang, and Yoav Goldberg. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. arXiv preprint arXiv: Arxiv- 2203.14688 2022b.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catan- zaro. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.13587, 2021.
James Thorph Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. Fnet: Mixing tokens with fourier transforms. arXiv preprint arXiv:2105.03822 2021.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895, 2020.
Bingyuan Liu, Yogesh Balaji, Lingzhou Xue, and Martin Renqiang Min. Analyzing attention mechanisms through lens of sample complexity and loss landscape, 2021a. URL https: //openzeview net
Hanxiao Liu, Zihang Dai, David So, and Quoe V. Le. Pay attention to mlps, 2021b.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. arXiv preprint arXiv: Arxiv-2202.05262. 2022.
Aaron Mishkin, Arda Sahiner, and Mert Pilanci. Fast convex optimization for two-layer relu net- works: Equivalent model classes and cone decompositions. arXiv preprint arXiv:2202.0133: 2022.
10

---

Under review as a conference paper at ICLR 2023
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.6614, 2014.
Mert Pilanci and Tolga Ergen. Neural networks convex regularizers: Exact polynomial-time convex optimization formulations for two-layer networks. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning. volume 119 of Proceedings of Machine Learning Research, pp. 7695-7705. PMLR, 13-18 Jul 2020. URL https: http:://proeeeeiee ngs ml press vv v119 lanni220e html.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitring small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188)99 2019.
Yongming Rao, Wenliang Zhao, Zheng Zhu, Jiwen Lu, and Jie Zhou. Global filter networks for image classification. Advances Neural Information Processing Systems, 34, 2021.
Saharon Rosset, Grzegorz Swirszez, Nathan Srebro, and Ji Zhu. LI regularization in infinite di- mensional feature spaces. In International Conference on Computational Learning Theory, pp- 544 558. Springer, 2007a.
Saharon Rosset, Grzegorz Swirszez, Nathan Srebro, and Ji Zhu. LI regularization in infinite di- mensional feature spaces. In International Conference Computational Learning Theory, pp. 544-558. Springer, 2007b.
Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Un- raveling attention via convex duality: Analysis and interpretations vision transformers. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 19050-19088. PMLR, 17-23 Jul 2022. URL https: /ttpp/eeeingg....////// ww...rrees...............emmmmm. / v1 62 / sahiner22a. html.
Pedro Savarese, Itay Evron, Daniel Soudry, and Nathan Srebro. How do infinite width bounded norm networks look in function space? In Conference on Learning Theory, pp. 2667-2690. PMLR, 2019.
Maurice Sion. On general minimax theorems. Pacific J. Math. 8(1) 1958. URL https: 1ttp//100444220
//projecteuclid. org:
Sho Takase, Shun Kiyono, Sosuke Kobayashi, and Jun Suzuki. On layer normalizations and residual connections in transformers. arXiv preprini arXiv:2206.000 00300000 2022.
Yuki Tatsunami and Masato Taki. Rafumlp: Do mlp- based models dream of winning over computer vision? arXiv preprint arXiv:2108.04384 2021.
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprini arXiv:2206.04811 2022.
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, et al. Mlp-mixer: An all-mlp architecture for vision. arXiv preprint arXiv:2 2021.
Hugo Touvron, Piotr Bojanowski, Mathilde Caron, Matthieu Cord, Alaaeldin El-Nouby, Edouard Grave, Gautier Izacard, Armand Joulin, Gabriel Synnaeve, Jakob Verbeek, al. Resmlp: Feedforward networks for image classification with data-efficient training. arXiv preprint arXiv:2105.03344, 2021.
Shikhar Vashishth. Shyam Upadhyay. Gaurav Singh Tomar, and Manaal Faruqui. Attention inter- pretability across nlp tasks, 2019. URL https://arxiiv abs / 909. 11218.

---

Under review as a conference paper at ICLR 2023
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov, Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:190S.09448. 2019.
Alex Wang. Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. 2019. In the Proceedings of ICLR.
Alex Warstadt, Amanpreet Singh, and Samuel R Bowman. Neural network acceptability judgments. arXiv preprint arXiv: 1805. 2018.
Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan. Metaformer is actually what you need for vision. arXiv preprint arXiv:2111.1144 2021.
Ming Yuan and Yi Lin. Model selection and estimation in regression with grouped variables. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 68 (8(1):49-67 67, 2006.
12

---

Under review as a conference paper at ICLR 2023
Appendix
Table of Contents
A Proofs of the results in the main paper
13
A.1 Proof of Lemma 1
13
A.2 Proof of Proposition
14
Proof Theorem
A.4 Strong duality proof
17
A.5 Proof of Theorem 2
17
A.6 Proof of Theorem 3
18
A.7 Matrix targets
20
A PROOFS OF THE RESULTS IN THE MAIN PAPER
A. PROOF OF LEMMA
We first note that similar scaling techniques were previously studied several papers, e.g., Lemma of Ergen & Pilanci (2021), Theorem of Neyshabur et al. (2014), Section 2 of Savarese et al. (2019), equation (2-3) of Pilanci & Ergen (2020).
We start with restating the optimization problem follows
h [Iwa) 1|2 + (waa))²
min EA
Σ
Σ w Wi2X:WayW3h.
+
(13)
WayER ER
j=1
We first apply the following scaling for (W2j. W3j )
W3j aj
W2 == a W2js W3j
(14)
where a, > 0. Since this scaling doesn't 't change the output of the network, i.e.,
w X.
=
X.
the training loss part of the objective function stays the same. Thus, we can directly search for the optimal scaling parameter a, > 0 by minimizing the regularization term via the following AM-GM inequality
(way)
h
Σ J=1
1| + (w0z))2
Σ
2
W2j
(1ww)il2 (wail) 2
> (Il (1 llz Wayl)
j=m
where the equality is achieved when aj the regularization term is in multiplicative form as follows
Thus, we obtain a reformulation of (13) where
W 1
w X, W2r Bi
3
| wayll
(15)
Way EA
WayER' May ER
Next, we apply variable change to the reformulation in (15) as follows
W2j
why =
w/y == U33 |[2 1
W2j
1
13

---

Under review as a conference paper at ICLR 2023
With this variable change, rewrite (15)
min EA
Ec c
Σ
w
Wijx,w'y
Wgg. Vi
+ 3
(16)
Wes
|=1
"ay
This concludes the proof and yields the following equivalent formulation of (16)
N
h
min w1y
Σe
L
Σ
w 1j X, W/y
Wajj Vi + 3 ||w/ll1-
EA
2=1
Way ER
due to the We also note that the equality constraint || و1g 1 can be relaxed || l1 S optimality conditions arising from the regularization term |/w3ll
A.2 PROOF OF PROPOSITION
We first note that in order to maintain strong duality in our convex problem derivations, we basically use the arguments in Pilanci Ergen (2020); Rosset al. (2007b), where the authors proved that as long as h exceeds certain threshold h", then there will be sparsity in the solution due to the sparsity-inducing regularization in (7). And we have the following upperbound h" S N + 1. Note that this N + upperbound is the worst case scenario and h- << N + in practice as validated in Pilanci & Ergen (2020). Thus, below, we assume that there is a sparsity pattern in the solution.
Given optimal solution to (8), denoted Z" E Raad first rewrite this solution as summation of rank- 1 matrices as follows
h
z' Σ 0, =
Zj
|zyl|2
j=1
where e, E R is the jth ordinary basis vector and we assume that there are h nonzero rows out of rows of Z due to the -induiiinn regularization in (8). Then, this implies that the output of the optimal can be equivalently formulated as follows
h
trace (2'x,) X,
Σ e
X
Zj
|21l2
Z
j=1
wi, = ejj W
Z,
V /|11112
V11z,112
wi X, w
where (wij, why W3j 1/ j=1 denotes an optimal solution to (6).
Next, we show that both of these solution sets achieve the same objective value
(1Xa : Σe c
wij
X. wzg. Ui
11= + (wzr) 2)
)²)
j=1
(1
Σ c
Σ c X;
2j
|1,
3
Σ
2j
(VIlzyla)
|22l122
2
| 2
j=l
N
Σ 3
X,Z), 3/ e T
+
(l|z, 11 + (|zz)112)
2
1=1
h
C
Σ trace (z,e} X.) ,
+ 3 Σ
|z ll2
j=1
C (trace (z xxxx)) x.).,00 .11) +B Σ
|zz 1
(17)
1=1
14

---

Under review as a conference paper at ICLR 2023
where the last inequality follows from the fact that there h nonzero out of n rows of Z due to the aparsity-inducing regularization in (8). Note that (17) and (8) are the same objectives evaluated at Z", which concludes the proof.
Extension to multidimensional outputs in Section 3.3: Here we show that the proof above can be straightforwardly extended to the multidimensional output case in Section 3.3.
Given an optimal solution to (10), denoted as Zi E Rnxd we first rewrite this solution as a summa- tion of rank-1 matrices as follows
Zi
Σ
Ctk
Lik
|
Then, this implies that the output of the optimal can be equivalently formulated as follows
trace Zi X3
Zik
|zull2
wi,
elk: w2j
Zik
W3y el
V
ZiR 2
1
Vilzal
where { (Wij. Ways 3 denotes an optimal solution to (9). Note that here index j [hc] instead of [h] in the scalar output case.
A.3 PROOF OF THEOREM 1
We first provide summary of our proof strategy. For the derivations of the convex formulation, we basically need to find the bidual form of (6), i.e., the dual of the dual problem. Thus, we start with taking the dual of (6). To avoid nonconvexity in the dual problem, we reformulate the dual constraint, which makes problem nonconvex, as a convex constraint. Therefore, we obtain a convex dual problem. Then, take the dual of the dual problem to get the bidual formulation of (6). Since convexify the dual problem, the bidual formulation also a convex problem. Therefore, we achieve an equivalent convex formulation of the original nonconvex training problem (6). We also note that a similar proof strategy was also used in Pilanci & Ergen (2020).
In order to take the dual of (7) (i.e. restated below for the convenience of the reader) need to form the Lagrangian function for the following optimization problem
min Wiy
Σc C
W.jXX.2, ys
||wlll
[waj L
J=1
To construct the Lagrangian function, we first introduce an additional variable ý E RRN as follows
min YER" WiEEA [|way $1, ER
c (D v1) + B||lssl
s.t.
j =
w W W2, Wzj. Vi E [n]] (18)
Now we can form the Lagrangian for (18)
L(v.y.w3)
L (Ur, y1) + B ||
Di Σ
w
X. W22
j=1
N
N
N
h
C ( v.)
Σ
viii + [/
Σ
1/2
w 11 W2
1=1
/=1
N
C Cir. v) +
Σ
v. + B |//l
Σ
u.
m.
X.2, W3)
Minimizing the Lagrangian L(-) yields the following dual problem of (6)
max -L" (v,y)
s.t.
J X, w2 5B,
(19)
VERN
wi
W1EA.
1
15

---

Under review as a conference paper at ICLR 2023
where -* (-) denotes the Fenchel congregate function of the original loss function (+) (Boyd & Vandenberghe, 2004), which is defined as follows
3 (v.y)= = max z v - L (z.
ZERA
In order to convexify the dual constraint, next find the maximizers of the dual constraint as follows
N W
max w; ".E (aalllallllS1
Σ
max
Σ W
X,
1=1
i=1
max
w1 EA
max
WIk
0g
max
(20)
kejn
where the upperbound is achieved when each w, has is vector of zeros except a single one located at the index of maximum norm of weighted tokens.
Based on the observation in (20), we can equivalently write the dual problem in (19) as follows
N
d' max -L" (v,y) VERV
S.t. max KEIN
S B
N
max - - (v.y)
s.t.
S B, Ak E [m].
(21)
VERN
Next, we form the Lagrangian for the dual problem (21)
N
!)
L(v.y,
+
An
Σ u.Xik
and the corresponding optimization problem can be written in terms of the Lagrangian as
N
min max L(v.y,A) = | - 1yl1 A
3
X VER
Then, introduce additional variables rs Rdd equivalently formulate the optimization problem above as
min max min veRe railire $1
L (v.y)
A
8
r
Due to Sion's 's minimax theorem (Sion, 1958), we can change the order the minimization and maxi- mization to obtain closed-form solutions for the maximization over the dual variable v. This yields the following problem
min X20
min
d x Vb
+ 3
XE-
S1
km
km
Next, we apply variable change as 2ς = AArk, then the problem above reduces to
N
min zy (
C
z xax,
2 3
Σ
A
Av
1=1
16

---

Under review as a conference paper at ICLR 2023
From the KKT conditions, now that A 2 global optimum. In partieular, if X |z 1l. then one can further minimize the objective function by reducing the A and therefore Ao would not be optimal. With this, the problem be reformulated as
12
min Pig
Σ c
Σ
Xik Di
+
1
which is the same formulation with (8) and therefore concludes the proof.
A.4 STRONG DUALITY PROOF
To get the bidual of (7), we first utilize semi-infinite duality theory as follows. We first compute the dual of (21) with respect to the dual parameter to have
(
Wiyx:wsy w2), yi [/w2ll.s
(22)
where || represents the total variation norm of the signed measure μ.. Remark that (22) is fiinite-dimensiona dimensional training problem such as the ones in Bach (2017). Also, notice that this problem is convex withe respect to the linear measure μ (Bach, 2017). Therefore, strong duality holds, i.c., d' = Pso where d' denotes the objective value of (21). In addition to this, although (22) an drinite-dimensione theorem (Rosset et al., 2007a). Therefore, (22) is equivalent to the following problem
problem, it has most N + heads at the optimum due to Caratheodory
N
h"
Poo
min EA [|was 11
Σ
X, W3j + wlll 1
(23)
1=1
where h1- S N + 1. We note that that provided that 2 (23) and (7) are the same problems,
which proves strong duality, i.e., p² P&c = d", where p' denotes the objective value of (7). 0
A.5 PROOF OF THEOREM 2
We first apply the scaling technique Lemma for {waj, Wajly=
W3/
2j == Oj Waj
Waj Or
Then, following the same steps in Lemma (9) can be equivalently formulated as
N
3
min
C
X.W.
waajll
(24)
Way 2/ Was ER
2
21
1=1
Next, we again construct the Lagrangian function by introducing an additional variable ý; RS, Vi E [N] as follows
min $.Ex" waywey EA [[wa S1
E(.).))) + Wayll
8:=
Σ
w
X2 W3j, Vi [M. (25)
Now we can form the Lagrangian for (25)
h
N
h yi Σ
((vi)) v, (war)
= C y.) + 3
Σ
Wasll +
Σ
vi
w x, w2j W3j
==1
j=1
1=1
j=1
c ():.):)) + Σ
v. 91 + +
Σ
|Wajlll
1. x, W. W3
1=11
j=1
1=1
N
h
N
(yr.))) +
Σ
v y; + 3
Σ
|Wayllr
Σ
w 1 X.W waj
1=1
j=1
17

---

Under review as a conference paper at ICLR 2023
Minimizing the Lagrangian L(-) yields the following dual problem of (9)
max VLER"
-L" ((vi) (
S.t.
max W
Σ v.wi W
S B
(26)
where L (-) denotes the Fenchel congregate function the original loss function & (.) (Boyd &
Vandenberghe, 2004), which is defined as follows
L" ((vi) (xr)-mn)
max trace (z - (Z,
ZERNxC
where V, Y Rxxx are the matrix representations for the set variables {vi, y; y=)==r characterize the optimal layer weight explicitly, we next find the maximizers of the dual constraint as follows
i=1' In order to
Σ 1 X W2
max WIEA lE(c)
Daw1 X
WI WIEA, (Eall/wall,S
n
max max WIEA [E(O)
Σ
max max Wi (E(0]
Wik
UyXak
M
max max [E(O] RE(M
(27)
where the upperbound is achieved when each w, has is a vector zeros except a single one located at the index of maximum norm of weighted tokens.
Based on the observation in (27), we can equivalently write the dual problem (26) as follows
max
-L" ([v:))
S.t.
GaXi
3 E [n], Al E [c]-
(28)
VERN
Then directly following the steps in the proof of Theorem problem
yields the following convex optimization
min Z, ER
c (trace (Z) X.) ya) 3
A.6 PROOF OF THEOREMS 3
We first apply the scaling technique in Lemma 1 for [W2j, Way]y=
W3/
= a, W2j Waj :=
Then, following the same steps Lemma (11) can be equivalently formulated
8
min Way waylll Wa
Σc
C
(wIjX, X22)) Wa
waalll
(29)
2
Next, we again construct the Lagrangian function by introducing an additional variable ýs e RS Vi [N] as follows
N
min 9.ER" wiy EA Me
c () +
/ 1| s.t. y. Σ
σ (wT,X.w.) way, Vi E [N].
j=1
(30)
18

---

Under review as a conference paper at ICLR 2023
Now can form the Lagrangian for (30)
h +
N
h
((vi)==1 v,, v.[way)
Σ c ()
3
wajjll +
vi
ýi Σ
(wi, x; w2y) war
j=1
N
N
h
N
2 (V
+ -
v. ys +
3
|/3yll ||
Σ
v
a (wij x2 was
1=1
j=1
j=1
N
Σ
c (v.,),) + Σ
v. y1
+
|/a3lll
(w2s) vi wa
j=1
Minimizing the Lagrangian L(-) yields the following dual problem of (11)
max -C" ((vi) (x.). V.ER
S.C. Way EA. Way 51
max
Σ via (w] X w2))
3 Vj E [b]
(31)
where L (-) denotes the Fenchel congregate function of the original loss function (.) (Boyd & Vandenberghe, 2004), which is defined as follows
L ([vi) N (x))
max Irace (z - L Y),
ZER
where V, Y E Nx are the matrix representations for the set of variables (v..).)
We next note that we utilize the gated ReLU nonlinearity introduced in Mishkin et al. (2022). Thus, the activations [w]] x22) can be expressed
σ (w1)),w2)) 1:y Ni, 11,, W2j
h where := := 1 [u 20} and here {uij. U2 are fixed vectors that can be randomly selected. For instance, common choice is U1j I,) and u2j ~ (O Ia). For the rest of the derivations, we use this equivalent formulation of the activation function.
In order to characterize the optimal layer weight explicitly, we next find the maximizers of the dual constraint as follows
N
N
mAX max
Σ
via (w] T, (2
max max max jE[n] [E|c WI
Σ
Un
w
X
jE[h] W.
1=1
max max max jE(A] WiJ lE(c]
Σ
Uat 1 W
2
N
S max max max jE(k] WIJEA IE(c]
Wiyk
Σ
0z 1
max max max JEIA lEle REjN
Ua 1
N
(32)
where the upperbound is achieved when each Wiy has is a vector of zeros except a single one located the index of maximum norm of weighted tokens.
Based on the observation in (32), we can equivalently write the dual problem in (31) as follows
max-
-L" ((vi) (M)) S.L.
Σ
va Xy
Xak
S S E [m], vl e [c]. Vj E [M]. (33)
VERN
2
Then directly following the steps in the proof of Theorem 2 yiellds the following convex optimization problem
N
h
min Z
ΣV
c
1 tytrace (zJ.x.) .9a
+ 3
| ll
1=1
1=1
19

---

Under review as a conference paper at ICLR 2023
A.7 MATRIX TARGETS
We now consider the following vector output attention based model training problem, where the targets vectors, i.e., Y; Raxe,
N
h Σ /=1
[| + (was)²).
min
W W, - Y,
+ 2 ==1
(34)
WiEE
2
Way ER amay ER
1=1
Then the corresponding dual problem is as follows
N
N V, WiX,X. W
max WERRIE
(-IIV.
s.t. max W.E. Wall.11
trace
Σ
S
1=1
(35)
order to characterize the optimal layer weight explicitly, we next find the maximizers of the dual constraint follows
max WrEA,/|Wal1g5
trace Σ V. WiX.W
max WICA
Σ V.
W
max max
Σ
Va
W
Xij
(36)
[E(d] JEI4
1=1
Based on the equivalent formulation in (36), the dual problem in (35) can be equivalently written as
max V.ERaxe 2
Y1 + +| I
S.L. max WTEA
Σ v Wixx 18g
S Vj E [d], VI E [c].
(37)
The rest of the derivations directly follows from the proof of Theorem 2 and yields the following result.
Theorem A.1. Based the characterization of the dual constraint in (36), the non-convex opti- mization problem (34) can be equivalently cast as the following convex optimization problem
(2) 2
Z
min (11) 3 ER x
Xij
Ya
+
1,00
1,00
|=1j=
(38)
Remark A.2. Instead of the non-convex formulation in (34), we can also start from the following formulation in this setting
8
min WAyAA
Σ W Wi,K.w2,w.j
Y,
+
Σ
[| 112 + ||w3)11))
(39)
2 ==1
WayER ER
which also yields the convex formulation in (38).
20


Reviews:
[{"review_id": "zVp7JHFtyd", "reviewer": "Reviewer_5uWH", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "The authors study a relaxation of the attention mechanism. For this relaxation (which is still a non-convex problem) they propose a convexification. For single output models, they show how to solve the original problem using the convex equivalent problem. ", "summary_of_the_review": "I have concerns about the presentation of the results (see first bullet in Weaknesses) as well as the justification of some of the theory claims (see second and third bullets). In view of these, I am afraid I have to recommend rejection.", "strength_and_weaknesses": "Strength:\n- the paper is in general easy to follow\n\nWeaknesses: \n- The paper's contributions (as listed on the title, abstract, intro) are misleading. The papers does *not* analyze the self-attention mechanism (eqn (2)) but something different (eqn (5)). (W_1 is independent of X_i which introduces something different from self-attention. It means that the authors are trying to convexify another problem, not self-attention layer.) \nI understand that (2) is challenging. I can also understand that potentially (5) is already challenging and of interest, but then (a) the wording and the claims have to be adjusted and (b) it would be needed to show that (5) is (practically) interesting/relevant (Again, in experiments, they compare their results with equation (8) which is something different from self-attention layer.)\n\n- In Prop. 1 it is assumed that there are h non-zero rows on Z. There is no guarantee for that and is unclear how the prop changes otherwise. No analogue of Prop. 1 is given for vector case.\n- In the proofs of the main theorems, I have a question about the use of Sion's min-max theorem. Take for example Thm1 where this step is explicitly stated in the proof. As far as I know, Sion's minimax theorem is valid provided at least one set is compact, but 3 equations below eqn (20), but sets appear unbounded. \n- Also in proof of Thm 1, some more details should be given  the last part of\nproof when they use KKT conditions.\n- Some typos in the paper - e.g., equation (2). \n- Question: In the experiments, how do you measure test/train error of the convexified problem? It suffices to know Z and you don't need to solve for w's (with an alike of Prop 1)?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "q7WH9vNBHx3", "reviewer": "Reviewer_Dm3W", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "The paper relaxes the softmax function to a linear operation with unit simplex constraints. Then for the self-attention-only network with a scalar or vector output, the paper shows that the nonconvex optimization problem could be cast as a convex one. It also demonstrates that the convexified problem has an implicit regularization mechanism, which promotes sparsity across tokens. For the self-attention network with an FCN layer, the paper shows that the nonconvex optimization problem with the gated ReLU activation is equivalent to a convex optimization problem. Experiments with the nonconvex and convex approaches are presented. Interestingly, the paper shows that the convex approach could mitigate the impact of the grokking phenomenon. ", "summary_of_the_review": "The paper proposes the convex optimization for self-attention-only network with and without FCN layers. The paper also reveals the sparsity-inducing regularization and advantage in mitigating the impact of the grokking phenomenon. However, it does not compare with the regularized training problems in (2) with the variants of SGD. Also, the interpretability of the learned W_1 does not present. So the effectiveness of the proposed approach is not clear. This is the main weakness of the paper.", "strength_and_weaknesses": "Strength:\nBy the observation that softmax function output a probability distribution, the paper presents an convex alternative to the self-attention mechanism and reformulated the training problem as convex optimization problem. The paper reveals that the convex optimization problem has the sparsity-inducing regularization. It also shows that the convex problem could mitigate the impact of the grokking phenomenon.\n\nWeaknesses:\nThere are three main comments below, which need to consider carefully.\nFirst, the relaxation of the softmax function is the fundamental assumption, i.e., the reformulation in Eq.(3). The relaxed linear operation does not depend on X_i, which may be not an appropriate approximation. Also, the interpretability of the approach does not explain clearly. Moreover, the paper does not verify the effectiveness of the relaxation. For example, it is better to compare with the regularized training problem in (1) using SGD.\nSecond, the shapes of the input and output of a Transformer block are usually the same. However, the paper only demonstrates the self-attention-only networks with the scalar and vector output.  What about the matrix output with the same shape of X_i?\nThird, the paper tests with deep settings by stack the convex transformer layers in the experiments. It is better to give more detailed explanations on this issue. Is there any difficulty to extend to deep Transformers?\nOther minor comments:\n1.\tThere are two “of “ letters at the second line of the first paragraph in page 2.\n2.\tIn (2), “X” should be “X_i”. \n3.\tAt the first line of the second paragraph in page 4, “varios”should be “various”.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "lq5EfoZUVq", "reviewer": "Reviewer_MHEr", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "This paper proposes a convex alternative to the attention operation in standard Transformers with the benefits of better optimization and interpretability. The analysis is insightful and deep, and can certainly enlighten further research in the community. ", "summary_of_the_review": "I think the experimental analysis is not sufficient. Although the theoretical analysis is nice, I still don't have enough confidence with the effectiveness of the method.", "strength_and_weaknesses": "[ Strength ]\n+ The motivation is clear and important. The interpretability of Transformers remains an open question in existing research, and it is important to the scientific understanding and design of each module. The paper follows the exploration line and steps into the problem from a new perspective, i.e., making the attention operation convex and easier to be optimized.\n+ The presentation is logical. It is easy to know the logic of the paper as it is very clear. Also, concrete proofs are provided, making the proposed method mathematically solid.\n+ Experimental results can well support the effectiveness of the proposed method.\n\n[ For further improvements ]\n- There is lack of complexity or efficiency analysis. \n- It would be more convincing if the authors can validate the proposed method on more downstream tasks. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "hYjsb76XB9", "reviewer": "Reviewer_qfuJ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "This paper derives a convex formulation for self-attention in transformers. The authors extend the formulation to various transformers architectures and then apply this convex formulation with an additional regularization for training transformers to find a globally optimal solution. The proposed formulation also suggests an implicit regularization mechanism induced by our attention mechanism. The authors verify the advantage of their convex formulation via fine-tuning experiments with BERT and algorithmic benchmarks.", "summary_of_the_review": "Overall, this paper could be an interesting algorithmic contribution. However, my main concern is the novelty of the paper. Also, more experimental results are needed to justify the advantage of the proposed convex formulation.\n\nCurrently, I am leaning toward rejecting the paper. However, given additional clarifications on the two main concerns above in an author response, I would be willing to increase the score.\n", "strength_and_weaknesses": "**Strong points:**\n\n1. The paper proposes an interesting convex formulation for self-attention. This convex formulation open up a new approach to studying attention mechanism and transformer. \n\n2. The paper is well-motivated and discusses relevant related works.\n\n3. The paper is well-written, and I really enjoy reading the paper.\n\n**Weak points:**\n\n1. My main concern is the proposed convex formulation is very similar to the convex formulation proposed in [1]. Even though [1] only discusses linear and ReLU activation multi-head self-attention, the techniques in both papers are quite the same. [1] also extends their results to other mixing mechanisms such as the MLP mixer and Fourier Neural Operator. \n\n2. My second concern is that the paper needs more experimental results on different benchmarks and applications to justify the advantage of the proposed convex formulation.\n\n**References:**\n\n[1] Sahiner, Arda, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. \"Unraveling Attention via Convex Duality: Analysis and Interpretations of Vision Transformers.\" ICML (2022).\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject