Attention as Inference via Fenchel Duality
Authors: Haoye Lu
Yongyi Mao
Maia Fraser
Venue: NeurIPS 2021 Submitted (2021)

Abstract: Attention has been widely adopted in many state-of-the-art deep learning models. While the significant performance improvements it brings have attracted great interest, attention is still poorly understood theoretically. This paper presents a new perspective to understand attention by showing that it can be seen as a solver of a family of estimation problems. In particular, we describe a convex optimization problem that arises in a family of estimation tasks commonly appearing in the design of deep learning models. Rather than directly solving the convex optimization problem, we solve its Fenchel dual and derive a closed-form approximation of the optimal solution. Remarkably, the solution gives a generalized attention structure, and its special case is equivalent to the popular dot-product attention adopted in transformer networks. We show that T5 transformer has implicitly adopted the general form of the solution by demonstrating that this expression unifies the word mask and the positional encoding functions. Finally, we discuss how the proposed attention structures can be integrated in practical models.

Attention as Inference via Fenchel Duality
Anonymous Author(s)
Affiliation Address email
Abstract
Attention has been widely adopted in many state-of-the-art deep learning models. While the significant performance improvements it brings have attracted great interest, attention is still poorly understood theoretically. This paper presents new perspective to understand attention by showing that it can be seen as a solver of a family of estimation problems. In particular, describe convex optimization problem that arises family of estimation tasks commonly appearing in the de- sign of deep learning models. Rather than directly solving the convex optimization problem, we solve its Fenchel dual and derive a closed-form approximation of the optimal solution. Remarkably, the solution gives a generalized attention structure, and its special case is equivalent to the popular dot-product attention adopted in transformer networks. We show that T5 transformer has implicitly adopted the general form of the solution by demonstrating that this expression unifies the word mask and the positional encoding functions. Finally, we discuss how the proposed attention structures can be integrated in practical models.
2
4
10
11
13
14
Introduction
15
16
Attention-basel deep neural networks are now integrated into cutting-edge language models that
have revolutionized broad range of tasks: machine translation 15 sentiment classification (77.
18
image captioning (29) and unsupervised representation learning 5 etc. Especially, attention plays
pivotal role in the construction of the transformer architecture (25), which has had profound impact
19
20
on the deep learning field.
21 Despite great empirical success, the driving principles of attention are still poorly understood. This
22 lack of understanding impedes practitioners from confidently and appropriately using attention layers and makes it challenging develop new attention based neural architectures.
24 In this paper, we offer new perspective for understanding attention by showing that it is in fact
25 solver for a certain type of optimization problem that corresponds to an inference task. We give
26 several examples, all of which can be characterized as follows: given 1) an unreliable estimate of
27 the mean of an unknown distribution p on Rd and 2) a preference distribution on Rd encoding
28 beliefs on p's selection, the inference task is to get a better estimate of p's mean given its unreliable estimate and We derive a convex optimization problem that is abstracted from the task and solve it improved
30 by instead solving its Fenchel dual 22 p. 104). Remarkably, the derived expression estimate of gives generalized attention structure whose special case is equivalent to the popular
32 dot product attention [15] that is also applied the transformer network [25]. In addition, we show that generalized attention expression has been implicitly adopted by T5 transformer [19] the expression unifies the concept of word masks and its positional encoding functions. Extra examples given show how the generalized attention structures can be used in practice.
34
Submitted to 35th Conference on Neural Information Processing Systems (NeurIPS 2021). Do not distribute.

---

36
2
Related Works
37
Since 2019, several authors have investigated the properties and working mechanism of attention. This series of works mainly addresses whether the attention mechanism can serve as proxy of 38
39
saliency 18 23 24
26 28]. Most of these works obtain insights into the attention mechanism by performing empirical studies. The related methods include analyzing the behaviours of trained attention-based networks [4, or pruning a few heads, or analyzing the effects of altering the attention
40
41
weights [8 26
mixture of these [9 24
Apart from understanding attention empirically, some theoretical results presented by Brunner el 44 al. [] and Hahn (7) show that the self-attention layers are not identifiable. This implies there could exist multiple combinations of attention weights that can provide equally good final predictions.
45
46 In particular, such non- -on-uniqueness means that the use of attention may complicate interpretability.
47
Another important approach to understand attention is to analyze its asymptotic behaviour when the number of heads and the network width approach infinity [8] 30). In this limiting case, the entire network can be seen as a Gaussian process [33 and its behaviours can be characterized by closed-fome expressions that are not available in the finite case.
48
49
50
51
Very recently (since 2021) several theoretical works have appeared that study attention outside the asymptotic regime. Lu et al. 14] set up a simple attention-based classification model and derive a closed-form relationship between the word's embedding norm and the product of its key and the query. They empirically show that such relationship also exists in a more complicated and practical configuration. Ramsauer et al. (20] construct an equivalence relationship between attention and
52
53
54
55
56
newly proposed Hopfield network with continuous states. In particular, they show that the new Hopfield network's update rule is equivalent the attention mechanism used in transformers 25
57
A Motivating Example
58
We first consider seemingly unrelated example, to illustrate the key ingredients of this paper.
60 Assume probability distribution p on Rdd has spherical Gaussian prior ~ (H, Ia). Let hp
61 denote the mean of the unknown p. Given an unreliable observation b of hp, what is the best guess 62 of hp? ? To solve this problem, we may formulate the following optimization problem
p' arg min P 2 b f ap(a) da
+ K(p, u),
(I)
with 0 responsible for the relative strength of the two terms, where K(p,, u) denotes the KL divergence between p and u. The basic idea behind (T) is that: although b not reliable, it should ap(a) da. Also, as encodes the preferred value of p. add the
64
65
not be too far from h, KL. divergence term to show preference for p that is close to u. As will be discussed later, such a formulation can be either obtained from the maximum likelihood principle or from the maximum entropy principle 11 In particular, Rioux et al (21) develop (T) for image de-blurring by applying Maximum Entropy on the Mean (MEM), information-theoreti method due Gamboa [6] but yet widely known in machine learning.
66
67
88
69
70
After obtaining the minimizer p" of (1). its mean f appp"a) (a) da gives our estimate of hp. Rioux et al. 21 prove, via Fenchel duality [2 p. 1.4] that the minimizer p" takes the form
72
u(a) exp(a, A+) (a) = J u((a)) exp(a', , Ar) dal
73 where
X arg max (b, A) XER
||a12 log u(a) exp(a, A) da.
2a
J
74. Note that fRx u(a) exp(a, )) da = exp (1, 2) + 2 ||]|3 as it is the moment generating function
75 (MGF) of u ~ [M, Ia). Substituting the expression into 3 followed by setting the derivative with
78 respect to À to zero yields X a+1 (b- 1). By 2 pa (a) x exp(-1 1 ||a ral2 + (a, X²)) x 2
77 cxp(-1 2 ||a (n 2)) Substituting X = o+1 (b 11) into it implies that p* follows a Gaussian
distribution NV W / + Ita La). Thus, our estimate of hp is Ital + r

---

79 In this paper, we focus on similar optimization problem that estimates hn assuming that a is instead discrete distribution. We show that such optimization problems naturally and frequently arise in 81 neural network designs. By solving the optimization problem, we derive closed-form approximation
80
82 for the estimate of hp. via Fenchel duality. The approximation then gives generalized attention layer
88 structure as shown in Fig [ A special case of it is equivalent to the familiar dot-product attention [TS]
84
that is also adopted in transformers [25]. Moreover, we will show that T5 transformer [19] implicitly adopts our generalized attention expression.
4 Setup of a Design Problem
86
Throughout the rest of the paper, we 88 consider a machine learning problem in which the objective is predict
(4)
(b)
c00
VIT
90 output quantity Y from a given input
y
X-
X. Additionally, Y may include K components, namely, be expressed as (y y y()).. (()). To be more concrete, we present a few example ma- chine learning problems and let them through development.
ged
ppt
92
93
85
218)
Generalized attention block
Probabitiies
2
Ei 20t wottd" aa (.
4.00
97 Example: Translation Problem. In this problem, the input X is a sentence, or a sequence of words, in the source language, and output Y is the sequence of words in the target sentence, where (k) denotes the j.th word.
Templates
(0)
supp.
99
Figure 1: A conceptual graph of the deep learning model that we work with. The block gkk) is the one we will investigate.
101
102
103
Example: Image Captioning. In this problem, the input X is raw image and output Y is the sequence of words in the caption, where (*) denotes the E word.
104
105
Example: Filling in the blanks task. This task has been used to train the BERT model [5]. The input X is a sequence of words with certain percentage of words masked. The output Y are the predicted masked words, where (k) denotes the t masked one.
106
107
108
The objective of any of these problems and that we address in this paper is to learn a function F, mapping from the space of X to the space Y so that = FXX). ). We will denote by f*() the part of responsible for predicting y(k) (Fig a),, namely, (k) F() (X). Although we here express F())), we note that it in fact possible that different F())ss
109
110
F as separate functions (F(1), F
111
share some component in common. We now focus on the design of (l))
112
We restrict the architecture of F*) to the form in Fig Ip with the main focus on the inference of (k), The extraction of feature (k) via two parallel modules r(() and prect that directly operate on the input X followed by function (8) (in Fig Ie), which we will design.
114
115
116 The Design Problem We describe the problem of designing g as follows.
117 Suppose that there is an unknown distribution (*)) on Rd whose mean vector is (*), namely,
(k)
(b) (a) da.
(4)
Rdd
Let (k) be another distribution on Rd that is generated as the output of network module ppal Here 119 (k) is referred to as the preference distribution, which serves as a prior guess of p(A),, Specifically
(8) puts non-zero probability masses on M "template" vectors {(4)) L22 ..... As in Rd, and their (k)
(kk
121 probabilities are respectively (4+)) ( ....,u "MM (which sum to 1). Collectively, we will denote the (k)) of templates by (4).
(k)
set [t. (k) to
The preference distribution (*) is considered good approximation of p (*), in the sense that the support of (A) is contained in the set (k) of templates. Note that if d is the word embedding space
3

---

word (banget)
zecovered At masked word
Softmon
Jont
Sofaaaxx
Add
FCN
AM4
A
Add & Norm
g
gk))
LSTM
ESTM
LSTM
Coabs supp poef wt.
(Larget)
Cuub. PipO. d pos wi.
smening
Tompless available be sclacaed
Saltsnol
Seaic
Transfure ation layees
CNN Layers
wied
word
word
word
word
(suros) (surcs)
(scanes)
Emnge
mascrd)
(a) Translation
(b) Image captioning
(c) Filling in the blanks
Figure 2: The model architectures of the three running examples. For the fevd in (a) and (b), the dashed links exist throughout the training and are replaced by the dotted ones the generation stage.
for a large vocabulary, and if the size M of the template set (kk) is relative small, then restricting the support of (8)) to within (kk imposes a strong constraint on p(A),
125
126
On the other hand, (k) is not a sufficiently accurate approximation of (A), in the sense that (k) may assign probabilities to (kk) somewhat differently. Such inaccuracy shifts the mean ((()) of (k) from
127
128
the mean h (*) of p(A), Suppose that there is another piece of information (( E Rd that is generated by another network module f(() and provides information regarding the mean shift. In particular, we assume that (A) a noisy version of the shift, more precisely,
130
131
(R) (k) - + 6,
(5)
where NN0,0 is the spherical Gaussian noise in Rd with covariance I We refer to 2()) as the evidence.
133
Then the design problem is construct function, or network block, g, which infers the un- known distribution (k)) and hence its mean l(k) based on the evidence (k)) and the preference distribution (k).
135
136
137
This formulation of the design problem might seem peculiar at the first glance, but we will show via examples (see 2 that such problem naturally arises in the construction of many machine learning models in practice.
139
Example: Translation Problem. For the translation problem, consider the model implementation plotted Fig 2a that is similar to the one proposed in [ We will focus on the part of the model responsible for inferring the tt word of the target sentence. In this model, (k) corresponds to the constructed feature according to ( that serves as an estimate of the context vector collecting the
141
142
143
source sentence information. The estimated (k(k) is then fed into a classifier y(k) to predict the tth
144
word. The preference distribution (*) is generated by p which takes the source sentence words as
145
inputs. particular, the support of (A) consists of the source sentence word embeddings (called annotations in [I] which are pre- pre-procesed by two LSTM layers The preference weight for each template depends on some positional encoding functions, which, principle, should assign higher
146
147
148
In this model, given input X, all d (A)ss share the same support T. The superscripts of the templates are then omitted to show their independence from k. Similar comments apply to implementations of the other two running examples.

---

weights to the templates appearing in the similar locations to the words we are inferring (that is, (A) assumed rely on the templates near t₁ more heavily).
149
150
Note that the inferred f support must be subset of "s as it is reasonable to assume that the target sentence words only depend on those appearing in the source sentence. Besides, although the preference weights specified by the positional encoding functions could provide some a priori information for the templates' weights in p( they cannot be accurate as their inferences do not consider the previously generated words ((itt) This results in the mean ye! shifted from (((k,, which is estimated by z(k) - that takes all the previously generated words y(f<)) into account using another LSTM layer. Thus, (A) and p(k) should not be far from 2 (k) + ((()) and ((,,,
152
153
155
156
157
158
respectively.
Example: Image Captioning. The caption generation model presented in Fig 2b has similar architecture reported in (29). This model shares the designs of f(() and f()) with the translation model while pped instead extracts the templates from a raw image using a CNN network. In general, word's position in the caption is independent of the location of the object it describes in the image. Therefore, in this model, all templates extracted by the CNN share the same preference weight.
160
162
163
As similar objects appear in image would have similar features extracted by the CNN (for example, a zebra and a horse), allowing similar templates not in T to participate in (*))ss ' estimation would possibly mix in information not contained in the raw image and harm the word inference accuracy. Therefore, we could improve the estimate of h (8) by choosing (k) similar to (A) in the sense that ( ' support cannot contain elements not in "s)
185
186
167
168
Intuitively, as the generation process proceeds, the context (k) should be updated to provide relevant information in the image to facilitate the next word inference. Such change is governed by the caption's 's semantic evolution, which is captured by (*)) = fevd that predicts the shift of the mean 4 (k) from (k), For this reason, j!lk) + ((k) serves as an estimate of (k) and should not be far away
169
170
171
172
from it. Likewise, x² should be close to p()).
173
Example: Filling in the blanks task. For the illing-in-the-bland tasks, let us consider model architecture plotted in Fig 2c that is similar to the one used in BERT (5). We focus on the inference of the hh masked word, which is assumed to be the jth word of the input sentence. In this model, f and (k) share the transformation layers (TL) that are commonly used in the NLP tasks to map one sequence of vector representations to another of the same length Taking the output sequence, f applies a linear map V to each of its elements form T as the support of (k) while the preference weights are specified by some positional encoding functions. At the same time, (*() = p(A) estimates (s'sss shift from the mean ((( due to the variation of the local information. For the same reasons discussed in the previous two examples, we need p.()) + z(k) close to (k) while ((k) is close to (kk)
175
176
177
178
179
180
181
182
Notably the formulation of the problem is based on the assumption that the network modules y!() and f are fixed and generate (²) and (k) satisfying the above assumed properties. In reality, (b) and r are in fact obtained via training. However, we argue that if g is made satisfy our 185
185
194
design objective, then we can at least interpret p( and ((8) obtained from training as serving to
186
produce 2(k) and (4) with our desired properties.
187
5
Formulation of an Optimization Problem
188
189
The discussion made in the previous section implies that the key optimization problem we are about to focus on should ensure
190
1. (b) is not too far from (R) + 2k), where (1) is constructed by (2) according to (4) and
(k) is the mean of the preference distribution u (*).
2. (k) is close to (A) while (t))s support must be subset of (A)'s. ss
Typical implementation of such layers include convolution layers, recurrent layers and self-attention layers.
5

---

194
These two desiderata prompt us to optimize:
a min 2 (1lt) + 2(h1) - ap(a) da + K(p, v14))
(6)
with a > O responsible for the relative strength of the two terms, where K(p, (()) denotes the KL divergence from to u (k), Remarkably, kp, (*)) has a finite value and only y has non-zero values on the support of u(()) Thus, both requirements in the second desideratum are satisfied by
185
196
197
using the KL. divergence as a measure for the closeness of p()) and u(*). Let p(k) be the minimizer of
198
The estimate / is
199
(x)
ap (a) da.
=
Rdd
200
Naturally, this optimization problem can be derived from two different, though, related perspectives
A maximum likelihond perspective. The optimization problem in (6) can be derived using the maximum log likelihood method by treating the KL-divergence term as regularizer. According to
202
the difference (p())+2(()) +2())) (k) follows Gaussian distribution (O, I This implies the log
203
likelihood function (k) - (]( + ((()) || Maximizing it with the KL-divergence
205
term as regularizer the same as minimizing
(1(t)) + =(m) (0) - n,p w(
(8)
20
206
where > 0 controls the strength of the regularization. Substituting 44 into 8 rearrangement yields
followed by
207
min
+ 2(8)) )
+2(k))
ap(a) da zd
2 K(pp u()))
2no02
which is equivalent to (6) by setting a no?
209
maximum entropy on the mean perspective Consider problem that seeks distribution p such
that the expectation JR ap(a) da not far from y(lk) 2(k), In particular, require
210
+200)))))) - ap(a) da
(10)
20
Note that, given 2x)), there are infinitely many p's that satisfy the constraints, which makes it difficult to pick "best" for later A technique known in information theory the maximum entropy the mean (MEM) 6 21 solves this problem by picking the best guess of the ground truth p" that
211
213
simultaneously satisfies 11) and minimizes the KL divergence to the preference distribution (*).
214
That is,
(10 + 2((0)) (x00))((00)
(K)) = arg min KC(p, u ())
subject to
ap(a) da S 2a
(11)
which is also the minimizer of 6 according Equation (18) of (1 and Corollary 4.9 of [2]
6
Optimal Solution
218 Rioux et al. proved that the optimization problem stated (6) has the following Fenchel dual (see
Theorem 2 of (2I)::
220
Theorem 1. The dual of (6) is given by
max
(2,11 2 201))
- |AAl - log
1
/ (a) exp(a, A) da
(12)
2a
Rd
221
Given maximizer x of 12
one can recover the minimizer ((k) of via
u(a) exp((a,>*
(k))
(13)
JA u(aa) exp(a' X+) da'
In fact, there is also Bayesian perspective derive the problem, which will be discussed in Appendix
A

---

By Theorem the estimated (b(4) defined in
7 can be re-written as
222
u ( (a) exp(a, X*)
(
a
F (a)
da-
a
Ja (a') expla' X²) da' da,
(14)
Rd
223
where X is a maximizer of (12)
In general, X does not have a closed-form expression in terms of a, (h) and (kk) and a standard paradigm is to search for it using gradient ascent-based methods. In this paper, we will not search for
224
225
X in this way: instead, we will derive a glosed-form expression to approximate it. Remarkably, this
226
227
takes the form of the generalized attention presented in [
Note that the integration in (12) equals E, (exp 2)]. the expectation of the random variable
228
(Xp X) where W has the probability distribution ((k) The expectation is just the moment
229
230
generating function (MGF) ()) of W, and the value log (() is called the cumulant of W 17
231
p.26], which has an expansion 17 (2.4)]
log (A) = ( A) + 2
(1)) )) + +(1AA1))
(15)
where (x( = Saw (a) da and (*)) = fR (a ( (a (()))) T ((a(a) da respectively
232
233
denote the expectation and the variance-covariance matrix of W.
234
Now we assume that a is small and we argue that this assumption is justified in practice. For instance, in the translation task, all of words in the dictionary can serve as candidate templates, which could be more than 10,000, but (k) reduces this size to the length of the source sentence (usually less than tens of words). The inference of p(() should strongly anchor around this prior information; consequently the information provided by z 2²() should weigh less. On the other hand, (*)) can hardly provide an accurate estimate of the mean shift, since the generation (R) is often ignorant of the templates selected by (A) (for example, in the example translation and image captioning models) or generated by a low-capacity module (as in the example filling-in-ll in-theebbank model). For these reasons, one should de de-emphasize the constraint imposed by 2k)) and hence choose a small a.
237
238
239
240
242
243
When a is picked to be small enough (see [22 ). the optimization of gets a large penalty on its L2 norm and thus, || is close to zero. Then, by 15 we have
244
1
log
u (k) (a) cap(a, > da log M(x") (1)) x²) (A", (()))
(16)
Rd
245
Substituting (16) into 12 followed by setting the derivative with respect to to zero yields
X = (La + -122)))
(17)
246
where La denotes the d x d identity matrix. As a is assumed close to zero, is further reduced to
x = az (M)
(18)
247
Plugging the expression into 14 gives the result stated as follows:
248 Theorem 2. For a small enough O, the estimated (k) defined can be approximated by
(k) (a) exp(a(aa (+)))
k((0)
a
Jga (()) exp(a(a', 2(()))))))))))))))))())))))))))) da'
da.
(19)
JRa
For the case that (((9) is discrete distribution with support {t(m), (k) , tt ...., (k) and the preference
(k)
250
probability {t ul ...., un (A))
(k)
}. 19 becomes simply
U expla(te,2()))
(R)
(20)
EE exp(a(tt,.(()))
In Fig we set d = 2 and visualize the approximation of (b) for different selections of a. We can
251
252 observe that, as a decreases, (20p outputs a better approximation of (k)
253 Let 1 we rewrite Theorem 2 as follows for later reference.
Corollary 1. For sufficiently large T, the best guess of h (k) defined in with = r equals
254
(1) u exp((t) (k) /(1))//))
(k)
(k)
(21)
Σラ u exp((t))) (8) 2/0))/ /T)
(K)
7

---

a 0.8
0.4
0.2
a= 0.1
1.0
(approc)
[appexx
papprox!
R Isporaxi
0.0
ulti 8.4
ult)) 0.
-0.5
otta 0.5
(00 -0.
0tb1 (akl-Ca 0LD
(90 0.5
-1.0
Figure 3: The approximation of (k) for different choices of a. The dots in light blue compose the support of discrete (k)) with the preference weights labelled above. The dark blue arrow starting from the mean ya()) of (*)) denotes the evidence 2 The red square marks the h(*) constructed by [4) with the X maximizes 12) while the green one marks the (a() approximated by (20). As
we can observe, 20 gives a precise approximation (8) when a is sufficiently small.
7 Discussion
255
In Section we derived an alternative expression of (k()) defined in by solving the Fenchel dual of the optimization problem stated in (6). Although the expression not in closed form, we are only
256
257
interested in the case when a is small. a closed-form approximation of (()) is derived in Theorem
258
and reduced to the form stated in (21) when considering discrete distribution (4))
259
As we pointed out, the block (*) in Fig 2a
Fig 2b and Fig 2c is expected to find the inferred p(*)
260
minimizing (6) followed by plugging it into 7 to construct (A))) Therefore, one can complete
the architecture designs of the three running examples by replacing g( with a network layer
282
263
implementing (1 namely, the structure in Figure (c).
264
The relationship between the optimal solution and the attention models. Remarkably, the expres- sion stated in (1) gives generalized attention block. By setting u; (k) for all i, the expression is equivalent to the well known dot-product attention [15), which is also applied in the transformer
286
network [25). The equivalence of the expression of (a))) and the dot-product attention layer tells us:
287
(a) by applying a dot-product attention layer in model, we essentially ask the model to perform
an optimization task defined in (6) and construct the output according to (7). (b) the derivation of (k) depends two relatively independent pieces of information: a preference distribution given the
269
271
global information and an estimate of the output's 's deviation from the preference distribution's 's mean according some local information. This suggests that the design of attention-based model can be decomposed into two parts that respectively estimate these two values.
272
273
The model consisting of stack of attention layers. Although our discussion focuses the case that contains a single attention layer, any attention layer in an attention stack fits our frameworks (see Fig D. In particular, all the attention layers closer to the input X than L can be grouped into the functions p(a)) or (k) For those layers that take the current layer's output as input, we can group
275
277
them into plle) where (k)) may contain the outputs of other attention layers working in parallel.
278
279
T5 transformer implicitly adopts the generalized attention structure. We now show that T5 transformer [19 can be seen as realization of the generalized attention in (2T), where the preference weights (k) unifies the concepts of word masks and T5's positional encoding functions. Consider the running example: filing in the blanks, with the preference distribution
280
281
282
10
if the 7th word is masked otherwise,
m ( (t,) = cxxp(b)i)/Z
(22)
where Z is normalizing constant and - is trainable scalar that only depends on the relative position of word and word j (which is the 4t masked word that we are inferring). Substituting such (4) into 21 with d yields
exp
(tr,)()))
+ -b + maskked (i)
1,000
t:
(23)
[tr.e/h12 + . + messkoo(!)
mml
Σn
exp
vi
8

---

286
where (i) is an indicator function that equals -oo if word is masked and zero otherwise. The expression in (33 has the same structure as that adopted in T5 transformer. where the indicator function serves as the mask function to prevent the model from assigning weights to the masked words. In this way, the concepts of word masks and the positional encoding functions are unified
287
289
by (k) in 22). Conversely, T5 transformer is realization of the generalized attention with the preference weights (8) specified in 22
280
281
232
Generalized attention structures suggested by the optimal solution. While T5 transformer has implicitly adopted the generalized attention, (2I) hints further generalizations could be made. For instance, in T5 transformer, the function outputting template's 's preference weights only considers the word masks and the word's relative positions. This function could be generalized to also consider the input sentence contexts, and the output weights encode the importance of each word before giving the local information stored in ²(,) The same idea could be applied to the image captioning example to replace the uniform preference weights. By adding a neural network taking the input image to generate non-uniform preference weights, we devise a mechanism to estimate the importance of each
294
296
297
239
part of the image before the caption generation. In this way, the newly added network collects global information from the image to propose a preference distribution, which could be updated locally
301
based on current generation stage encoded in 2())
302
Moreover, although we mainly focus on the case when (*) is discrete, we want to emphasize that the analysis performed in Section also covers continuous u (*). This hints that continuous attention mechanism could also be implemented, which might prove to be useful in some applications.
304
305
308
The approximations performed in 15 and (18) have implicitly assumed that random variable W following distribution (*)) has bounded moments. For a discrete (M)
Limitations and other comments.
307
308
with fixed support T {t,, 122 ... tn J. all the moments bounded and we can always pick a small enough a (or equivalently large enough scaling factor T in Cor to make the approximation meet
309
310
supplied by some neural layers, which could output templates of increasing norms as the training evolves. This
our requirements. A concern may arise as the support T in our running examples
311
312
problem could be alleviated by adding norm regularization or using normalized templates instead.
313
8
Conclusion
314
This paper presented a new perspective to understand the attention mechanism by showing that it can be treated as realizing solver of family of inference tasks. These tasks involve improving the noisy estimate of a distribution p's mean by preference distribution that encodes some beliefs of p's value. We have used three running examples with the typical model architectures to show that such tasks naturally exist in neural network design. We then abstracted a convex optimization problem from these tasks and derived a closed-form approximation of the optimal solution by solving the problem's 's Fenchel dual. We find that the closed-form approximation can be seen as a generalized attention layer and show that one of its special cases is equivalent to the dot- poo-product attention adopted in transformers. We further performed an analysis on the general form and showed that T5 transformer implicitly adopts the generalized attention structure with attention weights unifying the concepts of the word masks and the positional encoding functions.
316
318
320
322
323
324
325
This paper is the first work that presents a principled justification for the design of attention modules in neural networks. The generalized attention structure presented in this paper potentially opens a door to wide design space. For example, the preference weights need not be derived from the positional encoding functions; they could integrate a variety of information provided by other components of the network. Additionally, this research might have pointed to new ways to analyze the functioning of a neural network component, namely, via isolating the component from the complex network structure and asking: is there a "local problem" that is solved by the design of this component?
326
328
330
332
Potential negative societal impacts. This paper presents a new perspective to understand attention and derived a generalized attention structure. Our work is foundational, which we believe does not have direct negative societal impacts. Due to the very wide range of applications of attention, such as self-driving [12] and healthcare [16), our work may have unexpected negative impacts on these areas.
334

---

References
336
337
[1] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. In ICLR 2015, pages 1-15, 2014.
338
339
[2] J. M. Borwein and A. S. Lewis. Partially finite convex programming, part i: Quasi relative interiors and duality theory. Mathematical Programming, 57(1):15-48, 1992.
340
341
[3] Gino Brunner, Yang Liu, Damián Pascual, Oliver Richter, Massimiliano Ciaramita, and Roger Wattenhofer. On Identifiability in Transformers. In International Conference on Learning Representations (ICI.R), 2019.
342
[4] Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What Does BERT Look at? An Analysis of BERT's Attention. arXiv preprint 19Xiv:1906.04344 2019.
345
346 [5] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long and Short Papers), pages 4171-4186, Minneapolis. Minnesota, June 2019. Association for Computational Linguistics.
347
348
349
350
[6] Fabrice Gamboa. Methode du maximum d' d'entropie sur la moyenne et applications. Phd Thesis, 1989.
352
353
77] Michael Hahn. Theoretical Limitations of Self-Attention in Neural Sequence Models. Transac- fions of the Association for Computational Linguistics, 8:156-171, 2020.
355
[8] Jiri Hron, Yasaman Bahri, and Jascha Sohl-dickstein Roman. Infinite attention NNGP and NTK for deep attention networks. In International Conference on Machine Learning (ICML), 2020.
356
357
358
99] Sarthak Jain and Byron C. Wallace. Attention is not Explanation. Annual Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT), 2019.
380
361
[10] E. T. Jaynes. Information theory and statistical mechanics. Phys. Rev., 106:620-630, May 1957.
362
[11] E. T. Jaynes. Information theory and statistical mechanics. ii Phys. Rev., 108:171-190, Oct
363
1957.
364
[12] Jinkyu Kim and John Canny. Interpretable learning for self-driving cars by visualizing causal attention. In 2017 IEEE International Conference Computer Vision (ICCV), pages 2961- 2969, 2017.
366
367
[13] Jaehoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks gaussian processes. In International Conference on Learning Representations (ICLR), 2018.
368
389
370
[14] Haoye Lu, Yongyi Mao, and Amiya Nayak. On the dynamics of training attention models. In International Conference on Learning Representations, 2021.
371
372
[15] Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to iitention-basee neural machine translation. In Proceedings of the 2015 Conference on Empirical Methods in Nat- ural Language Processing, pages 1412-1421, Lisbon, Portugal, September 2015. Association for Computational Linguistics.
374
[16] Fenglong Ma, Radha Chitta, Jing Zhou, Quanzeng You, Tong Sun, and Jing Gao. Dipole: 376
377
Diagnosis prediction in healthcare via iitention-basce bidirectional recurrent neural networks. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD 17, page 1903-1911, New York, NY, USA, 2017. Association for Computing Machinery.
378
379
380
381
[17] McCullagh. Tensor Methods in Statistics: Monographs on Statistics and Applied Probubility. Chapman and Hal//CRC, Boca Raton, FL, first edition. edition, 1987.
10

---

383
[18]
Paul Michel, Omer Levy, and Graham Neubig, Are sixteen heads really better than one? In Neural Information Processing Systems (NIPS), volume 32. Curran Associates, Inc., 2019.
384
385
[19] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21 21110))1-67, 2020.
387
388
[20]
Hubert Ramsauer, Bernhard Schäft, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitner, Thomas Adler, David Kreil, Michael K Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. In International Conference on Leurning Representations, 2021.
390
381
392
[21] Gabriel Rioux, Rustum Choksi, Tim Hoheisel, Pierre Marechal, and Christopher Scarvelis. The maximum entropy on the mean method for image deblurring, Inverse Problems, oct 2020.
393
384
[22] R. Tymelll Rockafellar. Convex analysis. Princeton mathematical series; 28. Princeton University Press, Princeton, N.J, 1970.
385
396
[23] Sofia Serrano and Noah A. Smith. attention interpretable? Annual Meeting of the Association for Computational Linguistics (ACL), 2020.
[24] Shikhar Vashishth, Shyam Upadhyay, Gaurav Singh Tomar, and Manaal Faruqui. Attention Interpretability Across NLP Tasks. In International Conference on Learning Representations (ICLR), 2020.
399
400
401
[25]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. V. Luxburg, Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
402
403
404
405
[26] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head eeff-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, July 2019.
406
[27] Yequan Wang, Minlie Huang. Xiaoyan Zhu, and Li Zhao. Attention-based LSTM for aspect- level sentiment classification. In Empirical Methods in Natural Language Processing (EMNLP), pages
409
Austin, Texas, November 2016. Association for Computational Linguistics.
(28) Sarah Wiegreffe and Yuval Pinter. Attention is not not explanation. In 2019 Conference on Empirical Methods in Natural Language Processing and 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) 2019.
412
413
414
[29] K. Xu, J. Ba, R. Kiros, K. Cho, A. Courville, R. Salakhudinov, R. Zemel, and Y. Bengio. Show, Attend and Tell: Neural Image Caption Generation with Visual Attention. In Proc. 32nd Int. Conf. Mach. Learn., pages 257-261, 2015.
415
416
417
[30] Greg Yang. Tensor Programs I Wide Feedforward or Recurrent Neural Networks of Any Architecture are Gaussian Processes. In Neural Information Processing Systems (NIPS), 2019.
Checklist
419
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
422
423
(b) Did you describe the limitations of your work? [Yes] We have created a sub-section, called Limitations and other comments, in Section 7 to discuss the potential limita- tions of our work.
425
428
(c) Did you discuss any potential negative societal impacts of your work? [Yes] We have discussed this the end of the paper.
427
428
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
11

---

430
2. If you are including theoretical results...
431
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
432
(b) Did you include complete proofs of all theoretical results? [Yes]
438
3. If you ran experiments..
434
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or as URL)? [N/A] Our work presents a theoretical framework and does not contain any experimental study.
436
437
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they chosen)? [N/A]
438
439
(c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [N/A]
440
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [N/A]
442
443
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
445
(b) Did you mention the license of the assets? [N/A]
446
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
448
(d) Did you discuss whether and how consent was obtained from people whose data you'le 're using/curating? [N/A]
449
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
451
452
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
454
455
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
456
457
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
12


Reviews:
[{"review_id": "pSG392CICB", "reviewer": "Reviewer_p7xB", "summary": "This work constructs an optimization problem whose approximate solution results in a general form of attention. The optimization objective is derived from three assumptions: 1. the prior of the attention distribution has a support of the memory bank; 2. the likelihood of the query vector in the attention mechanism is roughly conforming to an isotropic Gaussian centered at the difference between the mean of the attention distribution and the mean of a prior distribution; 3. the standard deviation of the isotropic Gaussian is large. Then the mean of the posterior of the attention distribution is a general form of attention, with the dot attention scores reweighted by the prior probabilities. Under this framework, when the prior is a uniform distribution, dot attention is recovered; when the prior is defined by positions and word masks, the attention used in T5 is recovered.", "questions": "", "limitations": "", "rating": 6, "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": "Strengths:\n1. The math is elegant. Under the three simple assumptions, a general form of attention happens to be the mean of the posterior.\n2. The general form of attention makes sense intuitively: T5 can be understood as one special case when position information is used, and there are also other works that can be viewed as special cases such as applying a content selector before attention (https://arxiv.org/pdf/1808.10792.pdf).\n\nWeaknesses:\n1. My main concern is the lack of intuition. While general attention follows from an optimization problem derived from the three assumptions, this formalism doesn't provide any more justification. In fact, I think directly proposing the general attention is more intuitive than going through all the math, especially considering that we are not treating the attention variable as a random variable or use its distribution in any sense (such as by marginalizing it, rather than simply using its mean under this work's formalism).\n2. I don't think it's reasonable to assume the standard deviation of the isotropic Gaussian to be large ($\\alpha$ being small under this paper's notation). There are situations where we use a high $\\alpha$. For example, there are cases when we use a low \"temperature\" in attention to get more spiky or even one-hot distributions (such as in Gumbel-softmax https://arxiv.org/pdf/1611.01144.pdf). While this paper argues that the prior distribution reduces the support of the attention from tens of thousands of words to several dozens, I don't think \"the inference of p(k) should strongly anchor around this prior information\" (L237): no matter how large $\\alpha$ is, the optimal solution of Eq. 6 always has the same support as the prior, so using whatever $\\alpha$ is always anchoring around this prior information if we just want to have the same support.\n\nSuggestions:\n1. I think maybe it'd be more clear to get rid of the superscript $^{(k)}$ and simply mention that only a single step is considered in the case of predicting a sequence.\n2. The motivating example is on a continuous case while the focus of the rest of the paper is on a discrete case. I think maybe you don't really need this motivating example.\n3. The discussion section mentioned that the proposed approach might generalize to continuous cases. I think this paper would be stronger if you can make that work (and the motivating example would make more sense). \n\nOverall, this paper is well-written and the math is elegant. However, the intuition and motivation of the objective function are not very clear and I am not convinced that having this objective function provides a principled justification for attention. Therefore, I am leaning towards rejecting this paper unless I missed anything.\n\n==== Post Rebuttal ====\n\nMy concern about the assumption $\\alpha$ being small has been resolved: it makes sense that in practice (transformers) $\\alpha=\\frac{1}{\\sqrt{d}}$ is small. \n\nBesides, it's interesting to see the extra experiment that shows optimizing 6) has a similar BLEU score as using $\\alpha z^{(k)}$, and that $\\alpha z^{(k)}$ is close to the optimal solution.\n\nTherefore, I'm increasing my score to 6. That being said, I think this paper still can be further improved: for example, the BLEU score in the added experiment is quite low (I'd suggest using a smaller translation dataset such as IWSLT14 De-En), and I think it would be much more convincing if you can just use a trained transformer (with normal attention), and show that optimizing 6) gets a similar BLEU score. Lastly, I think this paper would be much stronger if you can add some experiments with a continuous version.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "8UNgrjGpHdB", "reviewer": "Reviewer_ZmjE", "summary": "The paper introduces a novel interpretation of the attention algorithm and corresponding theoretical discussions. The paper suggests that the \"(weighted) attention algorithm\" can be interpreted as an approximate solution of a density estimation problem, where we have a noisy observation of a mean and a reference distribution (or \"preference distribution\" in the paper). \n\nIn the proposed interpretation, the density estimation solves a proximal-type optimization problem. The penalty function is a reverse KL to a reference distribution, and the objective function is an L2-norm between the model's mean and the observed mean. Here, the authors emphasize treating the observed mean as a sum of the reference's mean and a noisy observation of a mean-shift.\nSince the optimization problem satisfies the strong duality, the optimization can be solved by its dual; however, it is still challenging to obtain its analytical solutions. Thus, in the dual problem, the authors introduce to approximate a term, a cumulant generating function of a reference distribution, by its second-order Taylor polynomial. The authors show that the closed-form solution of the resulting approximate problem is the dot-product attention, weighted by a reference distribution.\n\nFurthermore, the authors show that the proposed interpretation includes a well-known T5 transformer as a special case. Assume we aim at solving the proposed optimization problem to estimate a density of a categorical distribution, where the distribution is on a set of M vectors (i.e., M keys). Given some context information, a set of networks predicts a preference distribution and a query, which is the mean-shift. The weight on the L2-norm in the optimization problem is the inverse of the temperature parameter. Then the T5 transformer's attention can be obtained by the expected value of keys under the analytical solution of the approximate dual problem. Thus, the authors demonstrate that the closed-form solution of the approximate dual problem is a generalized attention algorithm. Under the proposed framework, the design of preference distributions is flexible; for example, one can choose different types of context information depending on tasks.\n\nIn experiments, the authors discuss how much the deviation of the approximate dual's solution from the ground truth depends on \\alpha.", "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": "Overall I enjoy reading the paper and believe the paper will be an essential contribution to the ML community. However, the paper should improve the quality and clarify.  \n  \n(Originality & Significance)  \nI found that the proposed framework can provide a significant contribution to the ML community. For example, the authors can check if reducing the approximation errors improves the performance of the attention mechanisms. In the proposed framework, the authors take two levels of approximations. The first is to take the second-order Taylor polynomial of the cumulant generating function (CGF) in the dual problem. The second is the approximation from Eq 17 to Eq 18. If we focus on the latter, we notice that one can directly plug Eq. 17 into Eq. 14 without using Eq. 18. The solution contains the covariance of the keys, and the resulting attention naturally takes into account the correlations between keys. It would be exciting to see if the attention with the covariance terms will improve the performance.   \n\nSimilarly, the authors can discuss the approximation errors of the second-order Taylor polynomial. It would be very interesting to see if the exact solution or higher-order polynomial will improve attention algorithms.\n\n\n(Quality & Clarity)  \nThe abstract and introduction section doesn't provide enough information about the paper and potential discussions. For example, it is difficult to guess what \"a certain type of optimization\" the authors refer to and in what context the authors introduce the generalization of the attention.\n\nIt is challenging to clarify how the proposed framework is supposed to look. I highly rely on the reference, G. Rioux et al., '20. I believe that the paper should be improved. For example, one may include 1) the motivation of solving the dual instead of the primal, or 2) what conditions described in the proposed method provide strong duality.\n\nWhile Section 3-4 provides example cases in attention mechanisms, it isn't easy to draw a connection from the descriptions to the proposed framework.\n\nIn my understanding, the connection of the proposed framework to the principle of maximum entropy is invalid. There must be a historical reason for naming \"maximum entropy on the mean (MEM).\" Still, it would be nice to clarify the difference between the MEM and maximum entropy principle. I also found that G. Rioux et al., '20. also clarifies that MEM is different from the principle of maximum entropy (see Sec 3.1 in their paper)  \n  \n   \n  \n(Minor comments)\n- In Eq 17, a \\alpha → 0, isn't \\Sigma^{-1} z the only one left?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "(Limitations)  \nIn my understanding, Theorem 2 is incomplete, while it is supposed to be the major contribution of the paper. In Theorem 2, the approximation error depends on the remainders of the Taylor approximation of the cumulant generating function (CGF) of the preference distribution. In understanding, the characteristics of the CGF's higher-order Taylor terms are not trivial, except the normal distributions. Therefore, I consider it essential to specify the conditions of the remainders and how small \\alpha should be to validate the approximation. One may only consider some common choices of preference distributions, such as the categorical or normal distributions. Note that the authors mention that in Lines 306-307, the approximations have implicitly assumed that the bounded moments of the preference distributions. However, I consider that Theorem 2 specify such conditions. Moreover, as far as I know, the behaviors of cumulants differ from moments in higher-order, so it is important to describe if preference distributions have bounded cumulants instead of bounded moments.\n\nWhile the paper provides a novel interpretation of the attention algorithms, the discussion about testing the hypothesis (= the proposed framework) is limited. As I mentioned in the \"Originality & Significance,\" one way to (indirectly) test if the proposed framework is generalized attention is to analyze if reducing the approximation errors improves the performance in a set of toy examples. I believe that providing such discussions will enormously improve the significance of the paper.\n  \n  \n  \n(Societal Impact)  \nN/A", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": ">12hrs"}, "extra_text": {}}, {"review_id": "5ZuIl0tOXB_", "reviewer": "Reviewer_4PEc", "summary": "The paper casts attention as a solver for an inference task that consists in estimating the mean of an unknown distribution given two information: an unreliable estimate and a prior distribution. The authors start by explaining this inference task and showing how this setup encompasses several real-life examples (translation, image captioning, filling the blanks). Then, they derive a convex optimization problem from this task and solve it by looking  at its Fenchel dual. Finally, the authors show that the solution obtained from this problem yields a generalized attention that is implicitly used in T5 Transformer.", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "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": "As mentioned to the AC, I would like to first signal that I am not familiar at all with attention and NLP tasks. Therefore, it is very hard for me to assess whether the model proposed in Section 4 by the authors is valid or not.  \n\nOriginality: The paper is definitely original and looks at a question that raises a lot of interest. The techniques used are classic in convex optimization but they are used in a novel way. I am not aware of any work modelling attention as a solution of a convex optimization problem. The related has the merit to cite other works that mathematically model attention. \n\nQuality: The paper seems technically sound on the optimization side. I also appreciate the fact that the authors bring the weaknesses and limitations of their approach. I would like to raise some points: \n\n1) The problem formulation (6) assumes that p is in the support of u^k (otherwise the KL divergence is not well defined). However, how do you enforce this fact? It does not seem obvious to me that you are enforcing this constraint in (6). \n\n2) Is it realistic in the case of language to assume that the noise in (5) is a Gaussian random variable? The whole analysis in Section 5 seems to rely on that.\n\n3) What is the choice of the prior that you would make when applying your method in practice? For instance in a translation task?\n\n4) Do you think that it is possible to numerically compare the Transformer T5 with a vanilla Transformer? How does your model differ from the vanilla attention?\n\nClarity: I think that the paper is very well written and I appreciate the efforts of the author to explain in detail their motivation (Section 3) and the setup they propose (Section 4) by taking concrete examples.  \n\nSignificance: I cannot really assess the significance of the paper as I am not an expert in NLP. I however believe that the authors should do more experiments on real-life NLP tasks and compare it to a vanilla transformer. This would help to assess the validity of their model and maybe figure out some limitations of their approach.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors dedicated a whole paragraph regarding the limitations of this paper which is definitely an advantage of the paper.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4 hours"}, "extra_text": {}}, {"review_id": "0DpVL9DRoo", "reviewer": "Reviewer_gF3T", "summary": "This paper provides a new perspective on the attention mechanism. More precisely, the paper claims that an attention module can be seen as estimating the distribution solving a specific optimization problem. This problem follows the assumption that many deep learning models use a module updating a global guess with local information. The problem consists in minimizing both the discrepancy between an adjusted value of the mean and the mean of the estimated distribution, and between a prior distribution and the estimated distribution. Using some approximation, the solution corresponds to a generalized form of the dot-product attention.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "Although the perspective provided seems at first glance artificial, it is original, sound, and seems indeed related to attention. The paper is globally well-written. To my mind, some points (how does this perspective apply in practice, some points to clarify) are worth being adressed before recommending acceptance. \n\nPros:\n- To the best of my knowledge, the paper indeed provides a new, nice perspective on attention.\n- The paper provides insightful links between their formulation of attention and classical deep learning architectures.\n\nCons:\n- The perspective proposed by the paper is potentially insightful but the authors does not bring out whether their perspective really applies in practice: would real, learned attention mechanisms indeed provide a good approximation to the optimization problem (6)?\n- Conversely, the paper could benefit from an example where solving the optimization problem can effectively replace a learned attention mechanism.\n- It is not clear to me whether this perspective really applies in transformers (see below). This may be due to a misunderstanding from my side though.\n\nQuestions and remarks:\n- l21-22: what is ``\"poorly understood\" in attention?\n- Section 2: There are works related to attention that could be worth citing in Section $2$, see non-local means [1], and attention as a kernel smoothing [2] which also proposes a generalized attention mechanism. \n- Section 4: previous work learn a template from the whole data set for computing the attention weights in linear time, e.g, [3][4]: would this improve the estimation of $p^{(k)}$ in your perspective? \n- l114: early explanation/intuition on $f_{prev}$ and $f_{evd}$ could be useful.\n- l120: inconsistency in fonts for the Real symbol.\n- l134:  could you elaborate on this problem? Why would we want to solve it in the context of deep learning?\n- l174: in BERT, the attention weights are computed with keys and queries matrices, and not with values and queries as suggested by Figure 2c. Moreover, the attention output is added to the input feature map which is different from what Figure 2c seems to suggest. What is your view on these two potential differences?\n- Figure 3: if $\\alpha$ is small, it seems that the estimate of $\\hat{h}^{(k)}$ goes away from $\\mu^{(k)} + z^{(k)}$, making the evidence term $z^{(k)}$ potentially useless?\n\n---------\n\n[1] Buades et al. : Non-local Means Denoising, IPOL 2011\n\n[2] Tsai et al. : Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel, EMNLP 2019\n\n[3] Mialon et al. : A Trainable Optimal Transport Embedding for Feature Aggregation and its Relationship to Attention, ICLR 2021\n\n[4] Kolouri et al. : Wasserstein Embedding for Graph Learning, ICLR 2021", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See above (Cons) for the limitations of the paper from my point of view. \nI think the authors correctly address the potential negative societal impact.", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper received an overall score of 6 and is therefore borderline.

All reviewers agreed that it's a very promising paper that could shed some light on attention models. 

However, in the current state, we believe the paper isn't ready for publication yet:

- The writing would benefit from some polishing. 
- Section 3 "Motivating example" feels disconnected from the rest of the paper. We would personally remove it since it appears in Theorem 1.
- Section 4, Figure 1 and Figure 2 were difficult to follow. Some effort needs to be spent on clarity.
- Reading the paper, we are left with the feeling "so what?". Although a variational perspective of attention is interesting, we already know empirically that attention works. Coming up with benefits that only this new view brings is very important. To that end, the experiments that the authors posted in the rebuttal are promising.
- Reviewer ZmJe mentioned issues with Theorem 2

Overall, we agree with reviewer ZmJe 

> In conclusion, the authors respond to most of my concerns. Unfortunately, however, such modifications require major revision, and thus I will keep my current score. On the contrary, I strongly believe that the revised version will be a significant contribution to the ML community.

We therefore believe that the ideas described in the paper are not yet clear enough to obtain the impact that they deserve. Polishing the paper and adding more experiments will greatly increase the impact of this paper and the authors will eventually benefit from it.
Decision: Reject