C ONVEXIFYING T RANSFORMERS : I MPROVING OPTIMIZATION AND UNDERSTANDING OF TRANSFORMER NETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁrst introduce a convex alternative to the self-attention mechanism and reformulate the regularized training problem of transformer networks with our alternative convex attention.
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.
I NTRODUCTION
Transformer networks proposed by Vaswani et al.
(2017) have become a dominant architecture in various tasks, especially Natural Language Processing (NLP)
(Devlin et al., 2018;
Radford et al., 2019), due to their extraordinary generalization properties and high capacity to learn from vast amount of data.
Although there exists substantial empirical evidence on the effectiveness of transformer networks, revealing the underlying theoretical reasons behind their success is still an open research problem due to their highly nonlinear and nonconvex structure.
A signiﬁcant body of research focused on analyzing certain components of transformer networks via empirical studies.
As an example, Liu et al.
(2021a)
; Vashishth et al.
Dong et al.
Voita et al.
; Takase et al.
; Liu et al.
(2021a) studied the impact of the attention mechanism on transformer networks.
Although these studies agreed that attention is an essential component 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 affecting the performance of the network, which is an indicator of large amount of redundancy in the network.
Vashishth et al.
(2019) provided a set of empirical evidence showing that attention might not be needed for some NLP tasks.
Additionally, Dong et al.
(2021) revealed that although attention is at the heart of transformer networks, training an 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 position of these might considerably impact the performance of a transformer network.
However, a 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
[Figure]
Tokens
Activation
Linear Layer
Tokens ≡ +
Regularization: Sparsity across tokens
[Figure]
Simplified Transformer with proposed attention
Our Convexified Transformer Block
[Figure]
Recently, another line of research has focused on understanding structures and patterns emerge throughout the training of transformer networks (Power et al., 2022;
Thilak et al., 2022; Barak et al., 2022).
In particular, the grokking phenomenon was ﬁrst observed by Power et al.
(2022) on speciﬁc algorithmic tasks, such as modular division operations.
Speciﬁcally, grokking refers to a sudden transition of validation or test accuracy to 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 deﬁnitely 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 ﬁrst propose a convex alternative to the self-attention mechanism and then develop our convex analytic framework on the resulting model as detailed in Figure 1.
1. C ONTRIBUTIONS
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/transformer networks with it.
• We convexify the regularized training problem of attention/transformer networks with the proposed attention layer as shown in Figure 1 and therefore enable ﬁnding a globally optimal solution without requiring any nonconvex optimization heuristic, e.g., layer normalization and skip connections.
• We also apply our 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 as a sparsity-inducing factor across tokens.
• We demonstrate the effectiveness of our convex reformulation via various experimental results.
We also show that our reformulation signiﬁcantly mitigates the grokking phenomenon studied in recent papers (Power et al., 2022; Thilak et al., 2022).
1. N OTATIONS
[Table]
We use lowercase and uppercase bold letters to denote vectors and matrices, respectively.
We denote a certain column/element of a vector or matrix using subscripts.
As an example, w jk denotes the jk th entry of the matrix W .
We use I k to denote the identity matrix of size k × k and(or 1 ) to denote a vector/matrix of zeros (or ones) with appropriate sizes.
We also use [ n ] for the set of integers ranging from 1 to n .
We represent the Euclidean and Frobenius norms as (cid:107) · (cid:107) 2 and (cid:107) · (cid:107) F , respectively.
We also use[ x ≥ 0] to denote the 0-1 valued indicator function.
We provide more notations we use throughout the paper in Table 1.
T RANSFORMER NETWORKS
Given a data sample (or sentence) X ∈ R h × d
as a sequence of h tokens with the embedding dimension d , we deﬁne the key, query, and value matrices as
[Equation]
which are the main components of the self-attention mechanism.
Then, a single transformer block, which is basically a stack of self attention, residual connection, layer normalization, and point-wise feedforward connections, can be formulated as follows
[Equation]
where σ ( · ) denotes the activation function for the FCN layer.
Although skip connections, layer normalization and FCN also play a crucial role in a transformer block, the success of these networks has been mostly attributed to the self-attention part, denoted as A o (Vaswani et al., 2017).
Therefore, in the following section, we ﬁrst study the training problem of a simpliﬁed transformer network, for which the network output is directly A o .
We then extend our derivations to a transformer network with FCN layers.
A TTENTION - ONLY NETWORKS
We ﬁrst consider a simpliﬁed transformer network only with a self attention layer that maps input sequence X ∈ R n × d to the output sequence ˆ Y ∈ R n × c with c outputs as follows
[Equation]
We also call the model (2) as an attention-only network.
This is a 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 a standard regression framework with an arbitrary convex loss function.
Given a training set { X i , Y i } N i =1 , where X i ∈ R n × d and Y i ∈ R n × c denote the input sequence and the labels/target outputs, respectively, the weight decay regularized training problem for the attentiononly network in (2) is as follows
[Equation]
where L ( · ) is an arbitrary convex loss function, including squared loss and cross entropy, and β >is the regularization coefﬁcient.
Although the attention-only model in (2) is quite powerful across various NLP tasks, e.g., natural language inference, neural machine translation, and text classiﬁcation (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 remedy these issues, in the following sections, we ﬁrst reformulate the training problem by replacing the attention part with an alternative convex layer and then cast the reformulated training problem as an interpretable convex optimization problem that enables the globally optimizing the network parameters.
3. C ONVEX ATTENTION LAYER
We ﬁrst note that since the softmax ( · ) operation is 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), we ﬁrst replace the softmax operation with a simpler yet effective alternative.
Particularly, since softmax converts the rows of its input matrix to a probability distribution, it can be relaxed as a linear operation with unit simplex constraints as follows for any U ∈ R n × n ,
[Equation]
where ∆ := { W ∈ R n × n : w i ≥ 0 , 1 (cid:62) w i = 1 , ∀ i ∈ [ n ] } denotes a convex set of constraints, also termed as unit simplex constraints.
Thus, we simpliﬁed and convexiﬁed the attention mechanism without disturbing its structure.
Based on this observation, (3) can be reformulated as follows
[Equation]
Note that the model above utilizes a single head attention model and, therefore, may not be practically relevant due to its insufﬁcient expressive power.
Thus, we introduce the concept of head to the problem in (4) as follows
[Equation]
Now, we are ready to apply the convex analytic tools to (5) as detailed in the next section.
3. C ONVEX OPTIMIZATION FOR ATTENTION - ONLY NETWORKS
As a warm-up, let us consider the scalar output prediction problem where the targets are onedimensional, i.e., y i ∈ R .
Then, (5) reduces to the following optimization problem
[Equation]
Next, we ﬁrst apply a rescaling between the parameters w 2 j and w 3 j such that (6) can be described as an (cid:96) 1 regularized optimization problem.
[Table]
Lemma 1.
The problem in (6) is equivalent to the following (cid:96) 1 regularized training problem
[Equation]
Based on 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 as the following convex optimization problem
[Equation]
Note that the equivalent convex model in (8) requires a single parameter matrix Z ∈ R n × d , where each row is the attentions scores of the corresponding token.
We also remark that the regularization in (8), i.e., the sum of (cid:96) 2 norms of the rows of the parameter matrix Z , is a speciﬁc type of regularization, also known as group (cid:96) 1 or Lasso, introduced by (Bakin et al., 1999) and shown to 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 in (8) as a 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 ﬁt the training labels { y i } Ni =1 .
Unlike the nonnegative attention scores in (6), denoted as w 1 j ∈ ∆ , the convex parameters Z ∈
R n × d 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, an optimal set of parameters for (6) can be recovered from a solution to (8) as proven in the following result.
Proposition 1.
After solving the convex optimization problem in (8) , one can recover an optimal solution to the nonconvex optimization problem in (6) , denoted as { w ∗ 1 j , w ∗ 2 j , w ∗ 3 j } hj =1 , as follows
[Equation]
where e j ∈ R n is the j th ordinary basis vector, z j ∈ R d is the j th row of Z , and we assume that there are h nonzero rows out of n rows of Z due to the sparsity-inducing regularization in (8) .
Proposition 1 proves that there is a one-to-one mapping between the parameters of the nonconvex formulation in (6) and the convex formulation in (8).
Therefore, there is no need to solve the challenging 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 1 to obtain an optimal solution to (6).
3. E XTENSION TO MULTIDIMENSIONAL OUTPUTS
In the previous section, we considered a setting with scalar target variables, i.e., y i ∈ R .
However, for some problems, e.g., multiclass classiﬁcation, target variables can be multidimensional.
Therefore, we now extend the analysis to the problems with multiple/vector outputs as follows
[Equation]
where y i ∈ R c and c denotes the number of outputs/classes.
Note that here we put (cid:96) 21 -norm on w 3 j 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 optimization problem
[Equation]
Theorem 2 shows that the equivalent convex model becomes separable over the output index l , i.e., instead of a single parameter matrix in (8), here we have c parameter matrices due to having c outputs in the nonconvex model (9) (see Table 2 for details).
This also illustrates that the number of outputs in the network directly controls the overparameterization level of the equivalent convex formulation.
3. A TTENTION NETWORKS WITH FCN LAYERS
Although the model in (5) exhibits interesting properties in various applications (Dong et al., 2021), it is basically a linear function of the token matrix X .
Therefore, it is likely to 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 conﬁrmed 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
[Equation]
where σ ( · ) is the activation function.
Theorem 3.
The nonconvex optimization problem (11) with the gated ReLU activation is equivalent the following convex optimization problem
[Equation]
where 1 ij := 1 (cid:8) u (cid:62) 1 j X i u 2 j ≥ 0 (cid:9) denotes the indicator function for gated ReLU activation and here { u 1 j , u 2 j } hj =1 are ﬁxed vectors that can be randomly selected.
Theorem 3 implies that introducing the activation function further increases in the overparameterization level of the equivalent convex formulation.
Precisely, (12) has h times more parameters than (10) as shown in Table 2.
N UMERICAL EXPERIMENTS
In this section, we present experimental results corroborating our theory in the previous sections.
Student-teacher setting with BERT: We ﬁrst consider a student-teacher setting with the pretrained
BERT model in the Hugging Face repository, i.e., bert-base-uncased .
Particularly, we feed the samples from the mrpc subset of the glue dataset (Warstadt et al., 2018; Wang et al., 2019) through the pretreained BERT model and save the input and output activations in a certain layer.
Then, we train the attention-only models, i.e., standard nonconvex self-attention (3), alternative nonconvex attention (9), and convex (10), from scratch using these pre and post activations as our training dataset.
All the experiments throughout this section are performed using a single GPU on
Google Colab.
We also use the same regularization coefﬁcient β and optimizer, i.e., Adam, and tune the learning rate and regularization coefﬁcient by performing a grid search on a validation dataset for both algorithms.
However, notice that we do not use any nonconvex optimization heuristics, e.g.,
[Figure] [Figure]
layer normalization and skip connections, for the convex model in all the experiments.
In 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 an order of magnitude smaller objective value than the standard nonconvex training, which is possibly stuck at a local minimum.
This effectiveness in training also translates into better generalization, i.e., our convex training approach obtains a lower test loss than the standard nonconvex training.
In order to understand the functions learned by each models, we also analyzed the attention maps in Figure 3.
Here, standard nonconvex training fails to learn the underlying model and outputs a uniform attention map across token.
However, our convex training outputs an attention map that is quite similar to 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.
Algorithmic datasets and Grokking: Inspired by the grokking phenomenon observed in Power et 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.
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 ﬁrst replicate the results in Power et al.
(2022) and conﬁrm
that the grokking phenomenon indeed emerges here, i.e., the nonconvex curve (purple) reaches 100% training accuracy at around 10 3 iterations in Figure 4a while it requires more than 10 5 iterations to 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 10 × faster than the nonconvex one in Figure 4b.
Moreover, the
[Figure]
cc
[Figure]
−
[Figure]
mod 97 .
Here, we train the networks to reach 99% test accuracy and show that our convex training approach exhibits a signiﬁcantly faster convergence and lower test loss than the nonconvex training.
[Figure]
Iterations 0.
[Figure] [Figure]
[Figure]
cc u r a c
[Figure]
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 a noticeable impact on the nonconvex formulation.
convex model also yield signiﬁcantly lower test loss in Figure 4c, which implies that it has higher conﬁdence in test predictions and therefore more robust than standard nonconvex training.
We remark that in the previous section, we theoretically analyze only single attention/transformer 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 to deeper settings.
We basically stack the convex transformer layers in (12) to obtain an arbitrarily 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 signiﬁcant improvements for the convex model, especially in terms of optimization speed, it fails to make any discernible difference for the nonconvex model.
Moreover, we run the algorithms on the modoperation which is basically more challenging task due to smaller number of samples.
In this case, one-layer models are not able learn the underlying task perfectly as demonstrated in Figure 6 but our convex model is signiﬁcantly 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 modcase, since both models fail to achieve perfect generalization in that case.
Figure 7a clearly shows that our 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 as demonstrated in Figure 7b, where we 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.
[Figure] [Figure]
# o f i t e r
[Figure]
we
6 since both algorithms fail to achieve 99% test accuracy.
demonstrate that the impact of grokking is substantially mitigated with convex
C ONCLUSION
In this paper, we studied the regularized training problem of attention/transformer networks and developed a convex analytic framework to train these networks.
Particularly, we ﬁrst proposed a convex alternative to the self-attention mechanism and then reformulated the training problem with this alternative attention mechanism as convex optimization problems.
Thanks to our convex reformulation, we globally optimize the network parameters without requiring any kind of nonconvex optimization heuristics.
In addition, the functions learned by our reformulation is transparent and interpretable.
More importantly, the reformulated problem reveals a sparsity-inducing regularization mechanism across tokens in the data, which also sheds more light on the structure of the resulting function and its generalization properties.
We then empirically veriﬁed effectiveness of our convex training approach over standard nonconvex training via several numerical experiments.
We also note that analyzing transformer networks through the lens of convex optimization theory is extremely crucial since it may result in substantial improvements in the understanding and optimization of these networks.
However, it is also quite challenging due to the inherent nonconvex structure of the network model.
To the best of our knowledge, this paper is the ﬁrst step in this direction and therefore has some limitations which can hopefully be eliminated by future work.
Speciﬁcally, in this paper, we mainly focused on the theory side of convex analysis and empirically validated the theory on a few small-scale problem instances.
We hope that a comprehensive and large-scale empirical veriﬁcation of our theory will be conducted by the follow-up papers. Table of
Contents
A Proofs of the results in the main paper
A.1 Proof of Lemma 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.2 Proof of Proposition 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.3 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.4 Strong duality proof . . A.5 Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.6 Proof of Theorem 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.7 Matrix targets . . .
A P ROOFS OF THE RESULTS IN THE MAIN PAPER A. P ROOF OF L EMMA
We ﬁrst note that similar scaling techniques were previously studied in several papers, e.g., Lemma 1 of Ergen & Pilanci (2021), Theorem 1 of Neyshabur et al.
, Section 2 of Savarese et al.
(2019), equation (2-3) of Pilanci & Ergen (2020).
We start with restating the optimization problem as follows
[Equation]
We ﬁrst apply the following scaling for { w [Equation]
[Equation]
where α j > 0 .
Since this scaling doesn’t change the output of the network
[Equation]
i [Equation] e., h (cid:88) [Equation]
[Equation]
the training loss part of the objective function stays the same.
Thus, we can directly search for the optimal scaling parameter α j > 0 by minimizing the regularization term via the following AM-GM inequality
[Equation]
where the equality is achieved when α j = (cid:113) | w 3 j | [Equation] 2 j [Equation]
Thus
[Equation]
we obtain a reformulation of (13) where the regularization term is in a multiplicative form as follows
[Equation]
Next, we apply a variable change to the reformulation in (15) as follows
[Equation]
With this variable change, we rewrite (15) as
[Equation]
This concludes the proof and yields the following equivalent formulation of (16)
[Equation]
We also note that the equality constraint (cid:13)(cid:13) w (cid:48)
2 j (cid:13)(cid:13) 2 = 1 can be relaxed as (cid:13)(cid:13)
w
[Equation]
(cid:13)(cid:13) 2 ≤ 1 due to the optimality conditions arising from the regularization term (cid:107) w (cid:48) 3 (cid:107) 1 .
A. P ROOF OF P ROPOSITION
We ﬁrst note that in order to maintain strong duality in our convex problem derivations, we basically use the arguments in Pilanci & Ergen (2020);
Rosset et al.
(2007b), where the authors proved that as long as h exceeds a 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 ∗ ≤ N + 1 .
Note that this N + 1 upperbound is the worst case scenario and h ∗ (cid:28) N + 1 in practice as validated in
Pilanci & Ergen (2020).
Thus, below, we assume that there is a sparsity pattern in the solution.
Given an optimal solution to (8), denoted as Z ∗ ∈ R n × d , we ﬁrst rewrite this solution as a summation of rank-1 matrices as follows
[Equation]
where e j ∈ R n is the j th ordinary basis vector and we assume that there are h nonzero rows out of n rows of Z due to the sparsity-inducing regularization in (8).
Then, this implies that the output of the optimal can be equivalently formulated as follows
[Equation]
where { w ∗ 1 j , w ∗ 2 j , w ∗ 3 j } hj =1 denotes an optimal solution to (6).
Next, we show that both of these solution sets achieve the same objective value f (cid:16) { X i , y i } Ni =1 (cid:17) :=
N (cid:88) i =
L   h (cid:88)
j = w ∗ (cid:62) 1 j X i w ∗ 2 j w ∗ 3 j , y i 
 + βh (cid:88)
j = (cid:16)(cid:13)(cid:13)
w ∗ 2 j (cid:13)(cid:13) 2 2 + (cid:0) w ∗ 3 j (cid:1) 2 (cid:17) =
N (cid:88) i =
L   h (cid:88)
j = e (cid:62) j X
i
z
j
(cid:112) (cid:107)
z j (cid:107)(cid:113) (cid:107) z j (cid:107) 2 , y i 
 + βh (cid:88)
j =   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
z j (cid:112) (cid:107) z
j (cid:107)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)  + (cid:16)(cid:113) (cid:107)
z j (cid:107) 2 (cid:17)

=
N (cid:88) i =
L   h (cid:88)
j = e (cid:62) j X
i
z
j , y i 
 + βh (cid:88)
j = (cid:0) (cid:107) z j (cid:107)
+ (cid:107)
z j (cid:107) 2 (cid:1) =
N (cid:88) i =
L   h (cid:88)
j = trace (cid:0) z
j e (cid:62) j X i (cid:1) , y
[Equation]
where the last inequality follows from the fact that there are h nonzero rows out of n rows of Z due to the sparsity-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 Z ∗ l ∈ R n × d , we ﬁrst rewrite this solution as a summation of rank-1 matrices as follows
[Equation]
Then, this implies that the output of the optimal can be equivalently formulated as follows
[Equation]
where { w ∗ 1 j , w ∗ 2 j , w ∗ 3 j } hcj =1 denotes an optimal solution to (9).
Note that here index j ∈ [ hc ] instead of j ∈
[ h ] in the scalar output case.
A. P ROOF OF T HEOREM
We ﬁrst provide a summary of our proof strategy.
For the derivations of the convex formulation, we basically need to ﬁnd 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 the problem nonconvex, as a convex constraint.
Therefore, we obtain a convex dual problem.
Then, we take the dual of the dual problem to get the bidual formulation of (6).
Since we convexify the dual problem, the bidual formulation is 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) we need to form the Lagrangian function for the following optimization problem
[Equation]
To construct the Lagrangian function, we ﬁrst introduce an additional variable ˆ y ∈ R N as follows
[Equation]
Now we can form the Lagrangian for (18) as
[Equation]
Minimizing the Lagrangian L ( · ) yields the following dual problem of (6)
[Equation]
where L ∗ ( · ) denotes the Fenchel congregate function of the original loss function L ( · ) (Boyd &
Vandenberghe, 2004), which is deﬁned as follows
[Equation]
In order to convexify the dual constraint
[Equation]
we next ﬁnd the maximizers of the dual constraint as follows
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where the upperbound is achieved when each w 1 has is a 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
[Equation]
[Equation]
Next, we form the Lagrangian for the dual problem (21)
[Equation]
and the corresponding optimization problem can be written in terms of the Lagrangian as
[Equation]
Then, we introduce additional variables r k ∈ R d to equivalently formulate the optimization problem above as
[Equation]
Due to Sion’s minimax theorem (Sion, 1958)
[Equation]
we can change the order the minimization and maximization to obtain closed-form solutions for the maximization over the dual variable v .
This yields the following problem
[Equation]
Next, we apply a variable change as z k :=
λ k r k , then the problem above reduces to
[Equation]
From the KKT conditions, we now that λ k = (cid:107) z k (cid:107) 2 at a global optimum.
In particular, if λ k > (cid:107) z k (cid:107) 2 , then one can further minimize the objective function by reducing the λ k and therefore λ k would not be optimal.
With this, the problem can be reformulated as
[Equation]
which is the same formulation with (8) and therefore concludes the proof.
A. S TRONG DUALITY PROOF
To get the bidual of (7), we ﬁrst utilize semi-inﬁnite duality theory as follows.
We ﬁrst compute the dual of (21) with respect to the dual parameter v to have
[Equation]
where (cid:107) µ (cid:107) TV represents the total variation norm of the signed measure µ .
Remark that (22) is an inﬁnite-dimensional 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.e., d ∗ = p ∗∞ where d ∗ denotes the objective value of (21).
In addition to this, although (22) is an inﬁnite-dimensional problem, it has at most N +1 heads at the optimum due to Caratheodory’s theorem (Rosset et al., 2007a).
Therefore, (22) is equivalent to the following problem
[Equation]
where h ∗ ≤ N + 1 .
We note that that provided that h ≥ h ∗ , (23) and (7) are the same problems, which proves strong duality,
i
[Equation]
e., p ∗ = p ∗∞ = d ∗ , where p ∗ denotes the objective value of (7)
[Equation]
A. P ROOF OF T HEOREM
We ﬁrst apply the scaling technique in Lemma 1 for { w 2 j , w 3 j } mj =
[Equation]
Then, following the same steps in Lemma [Equation] (9) can be equivalently formulated as
[Equation]
Next [Equation] we again construct the Lagrangian function by introducing an additional variable ˆ y i ∈ R c , ∀ i ∈
[ N ] as follows
[Equation]
h (cid:88)
[Equation]
w (cid:62)
[Equation]
[Equation]
Now we can form the Lagrangian for (25) as
L (cid:16) { v i } Ni =1 , y , { w 3 j } hj =1 (cid:17) :=
N (cid:88) i =
L (ˆ y i , y i ) + β h (cid:88)
j = (cid:107) w 3 j (cid:107) 1 +
N (cid:88) i = v (cid:62) i  
ˆ y i − h (cid:88)
j = w (cid:62) 1 j X i w 2 j w 3 j 
 =
N (cid:88) i =
L (ˆ y i , y i )
+
N (cid:88) i = v (cid:62)
i ˆ y
i + β h (cid:88)
j = (cid:107) w 3 j (cid:107) 1 −
N (cid:88) i = v (cid:62) i h (cid:88)
j = w (cid:62) 1 j X i w 2 j w 3 j =
N (cid:88) i =
L (ˆ y i , y i )
+
N (cid:88) i = v (cid:62)
i ˆ y
i + β h (cid:88)
j = (cid:107) w 3 j (cid:107) 1 − h (cid:88)
j
=
N (cid:88) i =
w (cid:62) 1 j X i w 2 j v (cid:62) i
w 3 j
Minimizing the Lagrangian L ( · ) yields the following dual problem of (9)
[Equation]
where L ∗ ( · ) denotes the Fenchel congregate function of the original loss function L ( · ) (Boyd &
Vandenberghe, 2004), which is deﬁned as follows
[Equation]
where V , Y ∈ R N × c are the matrix representations for the set of variables { v i , y i } Ni =1 .
In order to characterize the optimal layer weight explicitly, we next ﬁnd the maximizers of the dual constraint as follows
[Equation]
where the upperbound is achieved when each w 1 has is a vector of 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 in (26) as follows [Equation]
[Equation]
Then directly following the steps in the proof of Theorem 1 yields the following convex optimization problem
[Equation]
A. P ROOF OF T HEOREM
We ﬁrst apply the scaling technique in Lemma 1 for { w 2 j , w 3 j } mj =
[Equation]
Then, following the same steps in Lemma
[Equation]
(11) can be equivalently formulated as
[Equation]
Next
[Equation]
we again construct the Lagrangian function by introducing an additional variable ˆ y i ∈
R c
[Equation]
∀ i ∈
[ N ] as follows
[Equation]
Now we can form the Lagrangian for (30) as
[Equation]
Minimizing the Lagrangian L ( · ) yields the following dual problem of (11)
[Equation]
where L ∗ ( · ) denotes the Fenchel congregate function of the original loss function L ( · ) (Boyd &
Vandenberghe, 2004), which is deﬁned as follows
[Equation]
where V , Y ∈ R N × c are the matrix representations for the set of variables { v i , y i } Ni =1 .
We next note that we utilize the gated ReLU nonlinearity introduced in Mishkin et al.
Thus, the activations σ (cid:0) w (cid:62) 1 j X i w 2 j (cid:1) can be expressed as
[Equation]
where 1 ij := 1 (cid:8) u (cid:62) 1 j X i u 2 j ≥ 0 (cid:9) and here { u 1 j , u 2 j } hj =1 are ﬁxed vectors that can be randomly selected.
For instance, a common choice is u 1 j ∼ N ( 0 , I n ) and u 2 j ∼ N ( 0 , I d ) .
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 ﬁnd the maximizers of the dual constraint as follows
[Equation]
[Table]
[Equation]
where the upperbound is achieved when each w 1 j has is a vector of zeros except a single one located at the index of maximum norm of weighted tokens
[Equation]
Based on the observation in (32)
[Equation]
we can equivalently write the dual problem in (31) as follows
[Equation]
Then directly following the steps in the proof of Theorem 2 yields the following convex optimization problem
[Equation]
A. M ATRIX TARGETS
We now consider the following vector output attention based model training problem, where the targets are vectors,
i
[Equation]
e., Y i ∈ R n × c ,
[Equation]
Then the corresponding dual problem is as follows
[Equation]
In order to characterize the optimal layer weight explicitly
[Equation]
we next ﬁnd the maximizers of the dual constraint as follows
[Equation]
Based on the equivalent formulation in (36) [Equation] the dual problem in (35) can be equivalently written as max V i ∈ R n × c [Equation]
[Equation]
The rest of the derivations directly follows from the proof of Theorem 2 and yields the following result.
Theorem A.1.
Based on the characterization of the dual constraint in (36) , the non-convex optimization problem (34) can be equivalently cast as the following convex optimization problem
[Equation]
Remark A.2.
Instead of the non-convex formulation in (34)
[Equation]
we can also start from the following formulation in this setting
[Equation]
which also yields the convex formulation in (38)
[Equation]