LASER: Attention with Exponential Transformation
Authors: Sai Surya Duvvuri
Inderjit S Dhillon
Venue: ICML 2025 poster (2025)
Keywords: Large language modeling
deep learning
transformer
Conformer
ViT
Primary area: deep_learning->attention_mechanisms
TL;DR: We identify an issue with gradients backpropagating through standard attention mechanism in large language model and propose LASER to alleviate this issue.

Abstract: Transformers have had tremendous impact for several sequence related tasks, largely due to their ability to retrieve from any part of the sequence via softmax based dot-product attention. This mechanism plays a crucial role in Transformer's performance.  We analyze the gradients backpropagated through the softmax operation in the attention mechanism and observe that these gradients can often be small. This poor gradient signal backpropagation can lead to inefficient learning of parameters preceeding the attention operations. To this end, we introduce a new attention mechanism called LASER, which we analytically show to admit a larger gradient signal. We show that LASER attention can be implemented by making small modifications to existing attention implementations. We conduct experiments on autoregressive large language models (LLMs) with upto 7.7 billion parameters with an average improvement of upto 1.44% over standard attention on downstream evaluations and 1.65% finetuning improvements. Additionally, LASER demonstrates generalization performance improvement across a variety of tasks (vision, text and speech):Vision Transformer (ViT) on Imagenet, Conformer on the Librispeech speech-to-text and BERT with 2.2 billion parameters.

LASER: Attention with Exponential Transformation
Sai Surya Duvvuri' Inderjit S. Dhillon
Abstract
billoon parameter
Transformers have had tremendous impact for sev- eral sequence related tasks, largely due to their ability retrieve from any part of the sequence via softmax based dot-product attention. This mechanism plays crucial role in Transformer's performance. We analyze the gradients backprop- agated through the softmax operation in the atten- tion mechanism and observe that these gradients can often be small. This poor gradient signal backpropagation can lead inefficient learning of parameters preceeding the attention operations. To this end, introduce a new attention mecha- nism called LASER, which we analytically show to admit larger gradient signal. We show that LASER attention can be implemented by making small modifications to existing attention imple- mentations. We conduct experiments on autore- gressive large language models (LLMs) with upto 7.7 billion parameters with average improve- ment of upto 1.44% over standard attention on downstream evaluations and 1.65% finetuning im- provements. Additionally, LASER demonstrates generalization performance improvement across variety of tasks (vision, text and speech): Vision Transformer (ViT) on Imagenet, Conformer on the Librispeech speech-to-text and BERT with 2.2 billion parameters.
juoden (sttnii potesitti)
Dot-Produst Attenionn
Must the stannoon picbab fooo stud
postiaa 6 Leaing
LASER Allenton
Requine mininal cnde weficient attentioon pleesenbaiione
Figure I. Backpropagating gradients through the softmax opera- tion in the attention mechanism requires scaling with the Jacobian of softmax. We show that this Jacobian scales with attention probabiltiesswwightt which are typically small large language models (LLMs) with about 80% of the probabiltics less than 10 and about 20% less than 10-7 propose LASER attention that involves conducting dot-product attention with exp(-)- transformed value matrix V. i.e. conducting attention on exp(V). We show that LASER admits a larger Jacobian, is easier imple- ment and does not require any change to the underlying attention function, which may have Attentnt (Dao et al., 2022)). In the image, exp(-) and log(-) are element operations.
nuanced implementation (e.g..
Visicon
problem (Glorot & Bengio, 2010; Bengio et al., 1994). The attention mechanism plays key role in Transform- ers, where different weights or probabilities are assigned to token representations in a sequence, indicating their rel- ative importance, and these weights are computed via a softmax function (Vaswani et al., 2017). The Transformer architecture consists of multiple stacked layers comprising attention mechanism, where each layer operates the out- put of the previous one, forming the Transformer encoder or decoder. Learning within neural network is performed via gradient backpropagation, wherein gradients propagate backward through the network layer by layer using the chain rule (LeCun et al., 2002). During backpropagation, gradient magnitudes tend to diminish, resulting in weaker gradient signal reaching the bottom layers and inefficient learning, which is called as vanishing gradient problem (Glorot & Bengio, 2010; Bengio et al., 1994). Residual connections (He et al., 2016) are used in Transformers so that gradients can bypass the layers via skip connections during back- propagation improve the gradient magnitude for bottom
1. Introduction
Transformer architectures (Vaswani et al., 2017) have gained prominence over traditional models like LSTMs (Long- Sooot-Term-Memonee (Hochreiter & Schmidhuber, 1997) for various sequence-based tasks due their ability bet- ter capture long-range dependencies (Gemini, 2024; Meta- AI, 2024) without suffering from the vanishing gradient
'Department of Computer Science, University of Texas Austin "Gogle. Correspondence to: Sai Surya Duvvuri saisurya@ccsuuexxs..duu
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

LASER: Attention with Exponential Transformation
layers, reinforcing the idea that architectures capable of effi- cient gradient backpropagation tend to offer better training performance.
further. Efficient attention mechanisms have been an ac- tive area of research due to the quadratic computational complexity in sequence length of Attention, which prevents long mong-conext language modeling. One notable contribution is Linear Attention (Katharopoulos et al., 2020), which re- duces the quadratic complexity of self-attention to linear in sequence length by using kernel approximation of the softmax operation. Similarly, the Performer (Choromanski et al., 2021) develops alternative kernel approximation using random feature maps to achieve linear complexity.
In this paper, we theoretically analyze the gradient back- propagation in the attention mechanism of a Transformer and identify a vanishing gradient issue. During backprop- agation, the gradient can be scaled by very small value due the softmax operation in the attention mechanism. Based on this observation, we propose modification to the attention mechanism LASER LogArithm of Summed Exponentials of Representations. LASER is equivalent to conducting attention on exponentially transformed inputs and admits a log-sum-exp structure. We analytically show that gradients propagated via LASER attention are typically large. Since exp(-) transformation in LASER can lead to numerical overflow, develop novel implementation uog-Weighted-Sum-Be
A relevant work (Veličković et al., 2024) studies the impor- tance of expressing sharp distributions in attention mech- anisms of cansformer-bssee large language models. This helps with downstream tasks such as text retrieval. In con- trast, our paper finds that sharp distributions in softmax lead to gradient bottleneck during backpropagation and proposes LASER to fix this issue. In large vision transformer (Doso- vitskiy et al., 2021) pretraining, sharp softmax attention distributions led to training instabilities in (Dehghani et al., 2023), the authors propose oK-normalization, which applics layer normalization to queries and keys before computing attention weights. This equilibrates the norms of all the queries and keys, leading to uniform distributions. In (Kim et al., 2021), the authors identified that standard attention mechanism is not fz-Lipschitz and develop &g multi-head attention, which uses fz-divergence instead of dot-product between key and query vectors. This formulation admits a Lipschitz constant.
trick, inspired from the Log-Sum- Exp trick (Blanchard et al., 2019). This technique allows LASER to scale to large models with upto 2.2 billion pa- rameter models. We show that our implementation requires small modifications, and doesn't 't need any changes to the underlying attention mechanism which might admit a more nuanced implementation, for e.g., FlashAttentiin (Dao et al., 2022).
We conduct thorough empirical verification across a variety of Transformer models: Conformer (Gulati et al., 2020) for Librispeech speech-to- (Panayotov et al., 2015), Vision Transoor et al., 2021) for ImageNet classifi- cation (Deng et al., 2009), decoder-only text Transformer (Brown et al., 2020) C4 dataset (Raffel et al., 2020) and BERT ((eellin al., 2018)). We conduct experiments decoder-only autoregressive language models from 234 mil- lion parameters to 7.7 billion parameter models, where we demonstate improvements of up to 1.7% relative improve- ment in test loss over standard attention. We conduct one- shot evaluation on several downstream tasks and show that LASER outperforms standard attention with upto 1.44% ac- curacy improvement on average and 1.65% improvement average upon finetuning. On a 2.2 billion parameter BERT (Devlin et al., 2018), LASER gives a relative improvement of 0.93% on masked language modeling prediction error rate. LASER also demonstrates 1.2% improvement in accu- racy, and 0.2% improvement in validation word error rate in the Conformer benchmark.
3. LASER Attention: LogArithm of Summed Exponentials of Representations
We first formally introduce Transformers (Vaswani et al., 2017) and the underlying softmax dot-product attention in Section 3.1. In Section 3.2, we introduce LASER Attention by first deriving the gradients of standard attention by mak- ing observations on a simple case of sequence length 2, and then generalizing to larger sequence lengths.
3.1. Transformers and Softmax Dot-Product Attention
Let X E Nxd represent the input sequence with N tokens, where the i-th row is a d-dimensional representation of the i-th token. We descrihe the Transformer layer T: N + R- xd similar to (Katharopoulos et al., 2020) as follows:
2. Related Work
T(X) f(X attn(X)Wo):
(1)
The attention mechanism was used in Bahdanau et al. (2015) to drastically improve machine translation performance compared to encoder-decoder recurrent neural networks (RNNs) (Cho, 2014). This was later adopted in Transform- ers (Vaswani et al., 2017), which introduced self-attention to improve the performance in machine translation even
where f Rxd + R is usually implemented using 2-layer feed-forward neural network which acts on each token representation independently and Wo E dxd is a tunable parameter matrix. The attention function attn(.) is the only operation in the Transformer which is applied across the sequence axis. A single headed attention mecha-
2

---

LASER: Attention with Exponential Transformation
nism (Vaswani al., 2017) can be described as follows:
Dividing the numerators and denominators by exp(āsy)
gives:
Q =xWQ, K XWK; = XWv,
= QK
A=
a11 012
1 1+exp(i22 -021)
1+expla12 14oxp(@1=-á1) -m1)
expla11 -11))1--12]++ 812)+1
(t) (A)t,ax(AVV,
(2)
(22
0:2
exp(021 -
(a(ā1) ā12)
1-o(ā11 - -
(3)
where Q, K, V, attn(X) E Nxd The softmax (Bridle, 1990) operation is applied for each row ā of attention logits Ā = QK separately:
[ @22)
1 1- o -22)
where a denotes the sigmoid operation ( 1/(1 +
exp(-x))) As in (2), now multiply the attention prob- abilities A with the value matrix V. For simplicity, let the representation dimension d = 1, then the attention output will be as follows:
NV. Σ exp(ā,)
(softmax(ā)): exp(â.))
01 AV
Layer normalizations (Ba et al., 2016) are usually applied before or after f(-) and attn(-) (Xiong et al., 2020), but we omit this for brevity. A Transformer is a composition of multiple layers (1) - Tx.)), I f1..... L} sandwiched by the input embedding layer NVV + N and output softmax layer S Nxd NNV as follows:
Attention output: attu(X)
02)
(0((ā11 - u12)01 + (1 o (ā11 u12) )02
{( - 222)01 (1 o (Ō21
222) U02
(4)
where V = U1 We now find the gradient with respect
Transformer( (2) = S SoTo o................... Ti o (Z) ERV
to A via the chain rule:
ae
ae
i (attn(X)
where the inputs to the network 2 E R-V Thus choosing a suboptimal attention function attn(.) can affect every layer of the final Transformer()). Let ((Transformer(2) (Z), Y) be the loss function used to learn the parameters of the Transformer, where Y is the true label information. Autoregressive language modeling (Radford et al., 2018; Brown et al., 2020) involves using a causal mask M, which is a lower triangular matrix, and is added before the softmax operation as follows:
0A
① (X)
0A
hackpropggated gradient
Jacohian
Thus, small Jacobian magnitude can lead to small back- propagated gradient. We now analyze an element of attn Jacobian:
do,
= vāo - -12) o(ā - ā12))
0ā11
v20 - (12) oهā - x12))
attn(X) = softmax (M + QKTVV,
(v1 12) oân - 212 (ān - u12)) (5)
M,j = 0 if 2 else 00
possible saturation
The sigmoid function value, ققق x12) saturates to 1 when ă11 u12 becomes sufficiently large. Conversely, when āμ /12 is large and negative, the function value saturates to 0. In both cases, saturation leads to vanishing gradients, where the gradient becomes very small. This phenomenon is well-documented limitation of the sigmoid function (LeCun et al., 2002).
where @ denotes element-wise multiplication. During train- at D6 Dt
ing, the gradients bWx* BWa JWy are computed via back- propagation in a layer-by-layer fashion from layer L to layer and are used to update the parameters. In the next section, we analyze the gradient backpropagation through (.t and propose LASER attention.
3.2. Gradient Analysis of Attention
We extend this observation sequence length of size as follows:
For simplicity, we first let the sequence length N be 2 with attention probabilities = ") and attention logits as A QK Expanding the matrices A and Å, we get:
Lemma 3.1 (Gradient saturation in softmax). Let a E N be row in attention weights/probabiltiie A and similarly let à be a row in attention logits A, then:
forward pass:
a softmax((),
(111
012
ăn u12
A
softmax
(diag(a) - -aa" 0
021
0g2
@21
022
backward pass:
Dā
00'
exp(ān)
exp (xp()12) R12)
Da;
exp((11 texp( exp(
exp(a a (01) exp exp(22. (0x)(
softmax Jacobian:
a;(fii = j} 0.),
0ā,
exp(02) Hexp (12))
3

---

LASER: Attention with Exponential Transformation
where diag(a) denotes the diagonal matrix with diagonal elements
in a log-sum-exp form (Blanchard et al., 2019) follows:
01 = log(an exp(1i) + 012 exp(e))
We give proof of this lemma in Section A.2
log(exp(21 + log(an1)) + exp((2 + log(a12))) (8)
Key Observation. During the pretraining of 2.2 billion parameter autoregressive language model, we observe in Figure that about 80% of attention probabilities less than 10-3 and about 20% are less than 10-? From Lemma 3.1, can be seen that small attention probabilities d,. j {I....,Ny NN can lead to small Jacobian values, giving diminished backpropagated gradients.
Log-exp-sum function can be thought of as differentiable approximation of max function:
Lemma 3.2 ((Boyd & Vandenberghe, 2004)). The func- tion f(I1,..., = log (ea +.... e²n) is convex on R". This function can be interpreted as a differentiable approximation of the max function, since
m... S f(.1...... S f 21 II
To address this issue, now introduce LASER Attention which applies attention in exponential value space, exp(V),
+ log
as follows:
for all x E R". (The second inequality is tight when all components of x are equal.)
exp(laser(X)) = softmax(QK3 exp(V) laser(X) log(softmax((K") exp(V)) exp(()))
Given that max (.1,..,, In) function is not differentiable at points where two or more elements take the same value, log- sum-exp can serve differentiable approximation. Using Lemma 3.2, we can relate LASER (8) to max(-) operation as follows:
(6)
LASER Attention,
where log(-) and exp(.) are applied elementwise. Expand- ing (6) for N and = as done for standard attention (4) gives LASER output:
max (0) + log(an)) 02 log(a12))
S 01 S max(e) + log(an)) 02 log(a12)) log(2)
= ā12), 02 aāā2 - ā22) log((1 exp(r1) + (1 (x)
3.3. LASER Implementation via
01
=(oggos log((2 exp(e1) ) + (1 a2) expl
(7)
Wew--Wetihdd--m-----ee Trick
In this section we explore implementing LASER and pro- vide pseudocode. Given the log-sum-exp structure from (8):
Low gradient saturation. Computing an element in the Jacobian laser(X) /x) as done for standard attention in (5) gives the equation for LASER Jacobian:
Q1 = قāق - -12),
(exp(11) (xp(r)))a do
01 = log(a; exp(v1) + (1 (1)a1) a1) exp(e2)),
θān
an exp((1) + (1-a1)exp((2)
(exp(v))-exp(())) (xpp))))app(p) (e(p(()))) exp(p() ex)
notice that exp(.) operations can lead to overflow. This problem has been recognized in (Blanchard et al., 2019) and the "log-sum-exp trick" is used to avoid over- flows. However, the log-sum-exp trick cannot be applied directly as would difficult to implement without chang- ing the underlying attention function. We propose "log- weighted-sum trick", where subtract the maximum value max(u;, v2) from i, and vi and rewrite the above equation as follows:
We now consider limiting case to simplify the above equa- tion: exp(v1) >> exp(U2).
do1
a 01)
0ā
01 + exp((2)) (exp(11) exp(u2))
~ -1-a1
(1 قāق ā12)),
low saturation
01 = log((a: exp(v1 - m)
+ (1 - 01) exp(12 m)) exp(m))
where the approximation is due to exp(v2)/ exp(02)= xe O.
log((a exp(11 m) + (1 a1) exp(v2 m)) + m.
Relation between LASER Attention and max function. From definitions (3) and (7), LASER output can be written
Now conducting exp(.) operation on 01 - m and 022 - m will not lead to overflows. We can extend this to matrix-version

---

LASER: Attention with Exponential Transformation
(6) by conducting column-wise maximum of value matrix
Algorithm I LASER Attention with Log- Weighted-Sumi
V E Nxd follows:
Exp Trick
1: Input: Values V Rg Queries Q E &R Keys
K R
Mj = max Vg j {I,...,d} iE(looN)
Output: LASER Attention output O E R
Compute the column-wise maximum for the value trix V:
Define V E ER.v such that Vig- (Vi3 -m,)-
my sE(L....NG Vi j ft.......d
The above operations helps us conduct exp(.) operation without overflows. Then the final LASER attention opera- tion would be follows:
4: Subtract my from the jth column of V:
I/ Shift values to avoid overflow in the following
V Rxd such that V., - (Vi) m;)
0 log(softmax(QK") exp(() diag(exp(m))) ding(()iig()s))(m))
5: Apply attention with Queries Q, Keys and Values with mj. j (I....,d) added back the output, following (9)
Oū (logssotmax(((k))))o (xx()))()))))))))))) m,, (9)
0 R-Vxd is such that:
where O R.R. Nxd, Here, m = (mi...., ma) and diag(m) a diagonal matrix with elements of m as diagonals. Log- weighted-sum-exp trick allows us to implement LASER attention via merely modifying the inputs and outputs of standard attention, without changing the underlying atten- show in Section 4.1 that this
(0); [log(sottax(((K]] ((T))xx(()));; + m,
6: Return O
tion function. Additionally, trick helps avoid overflows in large scale settings. The fol- lowing JAX (Bradbury et al., 2018) code demonstrates that LASER attention can be implemented by utilizing standard attention functions.
4. Experimental Results
4.1. Autoregressive Language Modeling C4
In this section, compare the performance of LASER Attention with standard attention mechanisms in the context of an autoregressive language modeling task.
JAX attention
of LASER
Dataset and Setup. We use the C4 dataset (Raffel et al., 2020) for experiments. The training is conducted using batch size of 1024 sequences, each sequence has 1024 tokens. The models are trained for 160,000 iterations, re- sulting in the utilization of approximately 167.8 billion to- kens. Throughout the training process, we monitor both the training and test losses, and we observe a significant im- provement in the test set performance when using LASER Attention compared the standard attention mechanism (as illustrated in Figure 2). We use the AdamW optimizer (Loshchilov & Hutter, 2017) paired with cosine learning rate schedule (Loshchilov & Hutter, 2016) with linear learn- ing rate warmup followed by decay to zero at the end the training.
# given key (B, N, S), value (B, N, H, S), query H,
# B - batch size, N sequence length
H number of attention heads, S size of the head jap.max (value, axis=1, keep dims=True)
m = jax 1ax. stop.gradiinn (m) # stop the gradi ents al alongm
exp_value jnp.exp (value m) shift ing the values
= standarr # attention
implementation FlashAttention, etc.
Model Architecture. The base model architecture con- sists of 301 million parameters of decoder-only Trans- former, which is distributed across 32 layers as defined in (1). Each layer uses 8 attention heads, with each head hav- ing a size of 128. The MLP block this architecture, as defined (1), has hidden dimension of 2048.
at attention.out (key, query, exp_value)
out jnp. log (attention + m adding back the max values
5

---

LASER: Attention with Exponential Transformation
2.95
LASER attention
2.90
LASER attention Standard attention
0.25
Standard attention
0.24
2.85
grad norm
Irain Loss
0.22
2.80
2.75
0.18
2.70
2.65
20000
60020
80000 100000 120000 140000 160000 Step
2.60
20000
40000
60000
80000
100000
120000
140000
160000
Step
Figure 3. In this figure, we measure grad norm vs steps for an au- toregressive language model with 301 million parameters model corresponding to Figure
LASER attention
2,80
Standard attention
dient norm throughout training. An initial hypothesis was that higher gradient norms might lead more parameter change, consequently reducing the loss more effectively. To investigate this, we utilized the LAMB optimizer (You et al., 2019), which normalizes and renormalizes updates using the weight norm to ensure that the scale of updates matches the scale of the weights, thus voiding the effect of gradient/update norms on optimization. Interestingly, even with LAMB's normalization mechanism, we observed consistent improvement in training (Standard Attention: 2.749 vs LASER: 2. 2.736) and test loss (Standard Attention: 2.758 vs LASER: 2.741), suggesting that the performance gains were not solely driven by larger gradient magnitudes but are intrinsic to the model's 's architecture and the LASER Attention mechanism.
Test Loss
2.75
2.70
2.65
2.50
80000
80000
100000
120000
140000
160000
Step
Figure 2. Comparison of pretraining performance between LASER and the standard attention mechanism on a 301M parameter autore- gressive language model. The model consists of 32 layers, a 2048 hidden dimension, and a 1024 model dimension, trained on 168 hillion tokens from the C4 dataset. LASER attention consistently achieves lower training (top) and test (hottom) loss values.
Scaling to larger models. To demonstrate the scalability of our approach, we conducted experiments on a 1.1 billion and 2.2 billion parameter model. We note that without using the log- segg-weigted--ee sum--xp trick introduced in Section 3.3, the 2.2 billion parameter model training fails due to numer- ical overflow. In Table 2, we show that LASER attention outperforms standard-attention in a 2.2 billion parameter model with model dimension 2048 and hidden dimension 8192 with 32 layers and 8 attention heads (each of size 512). We also train a 1.1 billion model, which has a scaled down hidden dimension (4096) and attention head size (256). We show the training curves of both the models in Figure 5, Appendix A.3. In Figure 4, we used power law fit
In addition to this configuration, we also experiment with variant where the model retains 32 layers but increases the MIP block hidden dimension to 4096. In this variant, we increase the hidden dimension of the MLP block to shift more parameters into the MLP block. This configuration continues to show improvements both the training and test loss metrics, demonstrating that the effectiveness of LASER Attention is maintained even when attention parameters are reduced. The results of these experiments be seen in Table 1, where we also include ablation results showing improvements even with 16-layer setting.
Number of Layers
Hidden Dimension
LASER
Standard Attention
(() an to fit the final test loss values of L.L training runs as function of number of parameters. Additionally, we conducted training analysis, by qualitatively measuring stability of our training runs, sensitivity to bfloat16 training in Appendix A.5.
16
4096
2.673
2.681
32
2048
2.595
2,641
4096
2.555
2.575
Table 1. Comparison of test loss between LASER and Standard attention mechanisms across different distribution of parameters between MLP block f()) and attention attn(.) in (1), where we notice upto 1.74% relative improvement in loss.
Performance Analysis. We note that the 2.2B model with standard attention takes 27.22 hours on TPU v5 (Cloud, 2023) to reach a least test loss value of 2.327, however, LASER takes 24.88 hours (relative improvement of 9.4%)
Ablation with optimizers. For the 301 million parameter model, we noticed in Figure that LASER had higher gra-
6

---

LASER: Attention with Exponential Transformation
Scalling Law: Power-law fit
Evaluation on downstream tasks. In Table men- LAstR attention interpolated
tion the performance of our 2.2 billion parameter model several downstream tasks. Where evaluate on ARC (Clark et al., 2018), BoolQ (Clark et al., 2019), CB (Wang et al., 2019), COPA ( (Wang et al., 2019), HellaSwag (Zellers et al., 2019), MultiRC (Khashabi et al., 2018), Open- BookQA (Mihaylov et al., 2018), PIQA (Bisk et al., 2020), RACE (Lai al., 2017). ReCoRD (Zhang et al., 2018), RTE (Wang et al., 2019), StoryCloze (Mostafazadeh et al., 2016). WiC (Pilehvar & Camacho-Collados, 2019), Wino- grad (Levesque et al., 2012), Winogrande (Kocijan et al., 2020), and WSC (Wang et al., 2019). We found that LASER outperforms with upto 3.38% difference and 1% difference on average in accuracy.
Standard attention interpolated
LASER attention (cata)
202
Standard attention (data)
Test Loss.
2.50
2,43
2,43
2.59
2.30
0.55
0.50
0.75
1.00
3.55
1.5
3.75
2.00
3.15
Number of Parametess
Figure 4. Scaling law: Power-law fit for test loss against number of parameters. This plot uses 234M, 300M, 435M, 1.1B and 2.2B parameter models' final test losses after training on ~ 167B tokens. To reach a loss of 2.347, it takes 15.65% fewer parameters with LASER attention.
(a) Part 1
Dataset
+per-dim temp
LASER
+per-dim temp
Standard
+temp
WSC
77.89
79.65
Winogrande
62.27
61.72
62.35
60.69
62.19
Winograd
78.75
79.85
WIC
52.04
47.02
51.10
47.34
47.18
StoryCloze
77.87
77.77
76.48
76.22
76.17
Baselines Standard dot- not-product attention formulation in (Vaswani et al., 2017) uses temperature of Vd. Lower temperatures would be able express sharper distributions. Attention function with temperature hyperparameter can be
RTE
53.43
54.51
53.07
54.15
51.70
ReCoRD
83.29
85.05
85,10
84.82
84.89
RaceM
50.91
49.65
(b) Part 2
formalized as follows:
Dataset
+per-dim temp
LASER
+per-dim temp
Staudard
+4emp
QK
RaccH
37.88
37.94
37.65
37.62
37.04
Standard+temp:
softmax
V.
PIQA
77.09
77.30
76.88
76.55
76.45
rVd
OpenBookQA
48,80
47,80
48.48
QKT exp(v))
MultiRC
$7.03
57.78
53.94
55,36
53.38
HellaSwag
66.58
66.70
65.53
LASER+temp:
log softmax
Tva
COPA
$2.00
87.00
80.00
$2,00
85,00
CB
41.07
42.86
42.86
42.86
BoolQ
63.52
64.43
60.70
62,42
61.69
where the temperature hyperparameter is tuned manually. We now introduce per-dim temp operation, which scales each dimension with its own trainable temperature value:
Table 2. Accuracies of one-shot evaluation of a 2.2 billion parame- ter autoregressive language model trained via LASER and standard attention. We found that LASER outperforms or performs the same as standard atlention by up to 3.4%. On average, LASER gives an accuracy of 63.39% vs standard attention's 's 62.34%. Wc also include temperature-tuned version of standard attention, which gives an average accuracy of 62.49%. Additionally, we report per-dimension temperature scaling both LASER and standard attention, achieving overall means of 63.52% and 62.56% respec- tively.
@DDKT
Standard+per-dim temp: softmax
V,
Vd
D diag(softplus(p) = diag(log(1 exp(p))) e Raxd
where p is a d-dimensional trainable parameter and diag(softplus()) returns a diagonal matrix with positive diagonal entries to be scaled with queries. This technique is implemented in large language modeling frameworks such as PAX (Google, 2025) and AlgoPerf (Dahl et al., 2023), framework used to compare optimizers.
Finetuning on SuperGLUE. To further evaluate LASER, we fine-tuned the 2.2B parameter model on the SuperGLUE dataset (Wang et al., 2019) for 10 epochs. These experi- ments provide a more comprehensive view of how LASER improves over standard attention beyond 1-shot downstream evaluations. The results, presented in Table 3, show a 1.65% average improvement in decoding accuracy for LASER over standard attention.
In large vision transformers (Dosovitskiy et al., 2021), the norms of queries/keys can vary substantially across tokens with sharp softmax distributions in standard attention model training, leading to training instabilities. In (Dehghani et al., 2023), the authors propose OK-Normalization which con- ducts LayerNorm on queries and keys before computing attention weights:
Training and evaluation. All experiments are conducted using the PAX framework (Google, 2023) built on JAX (Bradbury et al., 2018), and executed on TPUv5 chips (Cloud, 2023). We use 64 chips for 300 million param-
Standard+QKK-Nommaiizttooo
softmax(LayerNorm()) LayerNomnk))) /VVv

---

LASER: Attention with Exponential Transformation
Method
Valid Error
Test Error
Train Error
Table SSpeeGG.UE fine-tuning results for the 2.2B parameter model (decoding accuracy %). LASER shows a 1.65% average improvement over Standard Attention.
Standard
0.2532
0.3749
0.1678
Standard temp
0.2513
0.3742
0.1668
Standard + QK normalization
0.2455
0.3644
0.1557
Task
LASER
Standard
Standard per-dim temp
0.2607
0,3821
0,1775
COPA
57.00
58,00
LASER
0.2417
0,3593
0,1527
LASER + temp
0.2427
0.3633
0.1538
WiC
56.64
53,92
LASER QK- Oo-oomaalization
0.2372
0.3588
0.1427
WSC
40.38
36.54
LASER per-dim temp
0.2475
0,3698
0.1573
RTE
22.02
20.94
Average
44.01
42.35
Table 4. Error rates for ViT-S/16 on Imagenet. We bold the lowest
emrors within Standard and LASER attention variants. LASER maintains upto 1.15% improvement over standard attention vari- ants.
eter model, 128 chips for 1.1 billion and 256 chips for 2.2 billion parameter model. Each training run takes upto 24 hours. We conducted hyperparameter search on 16-layer model mentioned in Table 1 with 15 hyperparameters us- ing search space mentioned in Table 9 and use the optimal hyperparameter for larger models.
Table also highlights the consistent improvement of LASER over standard attention across different configura- tions. Temperature tuning does not introduce major changes in performance for either attention mechanism, though lower temperature values tend to degrade results due sharper at- tention probability distributions. Adding QK "K-nommalization substantially enhances performance, indicating its robust- ness across attention mechanisms. On the otherhand, adding per-dim temperature negatively impacts performance, likely due to hyperparameter tuning being optimized without its presence.
4.2. Masked Language Modeling via BERT
In the experiments far, the focus was mainly on decoder- only models, to diversify our evaluation we now shift to encoder-only model- BERT (Devlin et al., 2018) trained via masked language modeling (as opposed to next token prediction in Section 4. We train a 2.2 billion parameter BERT on MLPerf training data which uses wikipedia arti- cles (ML.Commons). We used model dimension of 2048, hidden dimension 8192, number of attention heads 16, each of size 256. We get better error rate of masked lan- guage model predictions LASER 0.2125 vs Standard Attention 0.2145 (0.93% relative improvement). One can note that LASER shows more improvement in autoregres- sive language modeling compared to BERT
Conformer Librispeech Speech-to-Text. We also evaluate the performance of LASER attention on the Lib- rispeech Speech-to-Text dataset (Panayotov et al.. 2015) using the Conformer model (Gulati et al., 2020). Similar to the ViT experiments, we use the AlgoPerf benchmark and perform hyperparameter sweep across 50 configura- tions optimize standard attention. We pick the optimal them for 5 different random seeds (for initialization) and report the validation curves correspond- ing to median in Table 5 where we demonstrate a reduction in word error rate (WER) (0.0828 + 0.0808) when using LASER attention.
hyperparameters,
4.3. Vision Transformer (ViT) and Conformer- Speech-to-Tess Textt
Vision Transformer (ViT) on Imagenet-lk. In this sec- tion, we experiment with the Vision Transformer (ViT)
Method
Valid WER
Test WER
Train WER
Standard
0,083192
0.050876
0.032594
(Dosovitskiy et al., 2021) variant- S/16 on the Imagenet- 1k classification task (Deng al., 2009) which is a part of AlgoPerf benchmarks (Dahl et al., 2023) for optimizer com- parisons. These benchmarks are identically implemented in init2winit framework (Gilmer et al.. 2023), build on JAX, which we use for our experiments. A hyperparameter sweep was conducted over 50 configurations on NAdamW (Dozat, 2016), focusing on the search space defined Table 8. We selected the best-performing hyperparameter configuration based on validation performance for standard attention, ran it for 5 different random seeds (for initialization) and report the corresponding training runs in Table 4, where we show that LASER attention provides a 1.15% absolute improve- ment in error rate (25.32% + 24.17%), i.., a 44% relative improvement over standard attention.
Standard temp
0.08362
0.050398
0.032275
Standard QK Normalizzaioo
0.082865
0.049747
0.029574
LASER
0.082892
0.049786
0.030084
LASER temp
0.082137
0.049384
0.032822
LASER QK Normalizatonn
0.080837
0.049231
0.031283
Table 5. Word Error Rate (WER) results for Conformer model with different attention mechanisms and modifications. We bold the lowest WER values within Standard and LASER attention variants. LASER shows upto 0.2% improvements on word error rates.
In Table 5, we found adding QK-Normalization enhances performance for both LASER and standard attention. Tem- perature tuning does not significantly impact standard atten- tion but provides slight benefits to LASER. Since per-dim temperature is built in Conformer model architecture (Dahl
8

---

LASER: Attention with Exponential Transformation
al., 2023) we don't 't conduct ablation study its effect.
original DiffTransformer paper, d, 2dg for the combined V), À is learnable scalar, and dg is the dimension of the key/query vectors per head.
Ablations in Large Language Modeling
We propose a modification, LASER+Diff by equipping the softmax dot-product attention terms with LASER modifications. The LASER attention mechanism involves conducting attention on exponentially transformed inputs, exp(V), and takes the logarithm of the result. Ap- plying this to DifTTransformer, get.
In this section, we apply LASER and observe how it per- forms under the following changes: a) increase the model parameters to test the scalability of LASER b) using a Diff Transformer (Ye al., 2024).
Scalability of LASER
Q1K7 Vdk Q2K3
To test the scalability of LASER, we trained a model with 7.7B parameters for 44B tokens. Compared the 2.2B parameter model (which had model dimension 2048, hid- den dimension 8192, 32 layers, and 8 attention heads each of size 512), the 7.7B parameter model was scaled along all dimensions: model dimension 3440, hidden dimension 1 11584, number of heads 16, and dimension per head 720. The downstream task evaluation for the 7.7B parameter models presented in Table 6.
DiffAttDLASER (X) = log softmax
exp(V)
- Aloo
softmax
exp(v))
Vdl.
This formulation aims to address potential backpropagation challenges due to the softmax operations in both attention maps of the Drf.Transfrrmrr. We trained a 2.2B parameter DiffTransformer with Model dim: 2048, Hidden dim: 8192, Number of attention heads: 8, Head size: 512, and Number of layers: 32. The models were trained 24 billion tokens. The results shown in Table 7.
Table Downstream task accuracies for 7.7B parameter models trained 44B tokens. LASER shows an average improvement of 1.44% over standard attention. More pronounced differences are observed tasks like BoolQ (+6%), CB (+1.78%), OpenbookQA (+1.8%), and RTE (+3.61%)
Table 7. Downstream task performance for 2.2B parameter Diff- Transformer and DifTTransfrreerLLASERR models, trained on 24 billion tokens. On average, Dift+LASER shows improvement of approximately 1% over Tanffrranfrmmrr.
Dataset
LASER
Standard
ArcE
52.48
52.69
BoolQ
62.45
56.48
Dataset
Diff-LASER DiffTransformer
CB
44.64
42.86
ArdE
49,28
49.20
HellaSwag
$7.16
56.02
CB
42.85
41.07
MultiRC
56.00
55.59
HellaSwag
51.93
51,58
MultiCC
55,13
52.82
OpenbookQA
45.40
43.60
OpenbookQA
44,20
43,00
RaceM
44.64
44.15
RaceM
42.40
40.73
RTE
53.43
49.82
RTE
52.34
50.90
StoryCloze
71.78
71.51
StoryCloze
71.03
71.29
WiC
47.34
47.34
WIC
50.00
49.53
Winogrande
58.41
57.77
Winogrande
56.04
55,80
Average
53.97
52.53
Average
51.52
50.59
5. Conclusion
4.4. Ablation with Diff Transformer
The Differential Transformer (DifTTransformer) (Ye al., 2024) introduces differential attention mechanism. This mechanism calculates attention scores as the difference be- tween two separate softmax attention maps to cancel noise and promote sparse attention patterns (Ye et al., 2024). The Diff Transformer attention mechanism is formulated as:
We identified bottleneck in the gradient backpropagation of attention mechanism where the gradients scaled by small Jacobian values while passing through the softmax operation. We alleviate this issue by transforming the in- puts and outputs of attention mechanism, and show that this leads to larger Jacobians the limiting case. We demon- strate the improvements in training performance over four types of Transformers spanning different modalities (text, speech and vision): (a) decoder-only (via Large Language model) upto 2.2 billion parameters, (b) vision Transformers on Imagenet, (c) Conformer on Librispeech speech-to-text and (d) encoder-only (BERT) with 2.2 billion parameters., where we show significant and consistent improvements in performance.
QiKt"
DiffAutn(X X) softmax
Vdl
Q2K3 Vdl
Xsoftmax
where Q1, Q2 K1, K2 E Nxd= are projected query and key vectors, V Nxd i is the projected value vector (in the

---

LASER: Attention with Exponential Transformation
Acknowledgements
Language models are few-shot learners, arXiv preprint 2tti::0005.440033 2020.
We thank Devvrit Khatri, Nilesh Gupta and Aditya Kusupati for valuable suggestions.
Cho, K. Learning phrase representations using mn encoder- decoder for statistical machine translation. arXiv preprint urXiv:1406. 1078, 2014.
Impact Statement
Choromanski, K., Likhosherstov, V., Dohan, D., Song. X., Gane, A., Sarlos, T., Hawkins, P., Davis, J. Q., Mohiuddin, A., Kaiser, E., et al. Rethinking attention with performers. In International Conference on Learning Representations, 2021.
The paper presents a technique which improves pretraining of transformers, a widely used architecture in machine learn- ing. There are many positive potential societal consequences of our work, none which we feel must be specifically high- lighted here.
Clark, C.,, Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, 2019.
References
Ba, J. L.,, Kiros, R., and Hinton, E. Layer normalization. arXiv preprint arXiv:1607.0644 2016.
Bahdanau, D., Cho, K., and Bengio, Y. Neural machine translation by jointly learning align and translate. In 3rd International Conference on Learning Representations, ICLR 2015, 2015, URL http://:: aaaxiviiiv..///.i///aaaa/aaaavva org/ abs/ 1409. 0473.
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C.,, and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 0545..5557.05 2018
Bengio, Y., Simard, P., and Frasconi, P. Learning long-term dependencies with gradient descent is difficult. IEEE Transactions on Neural Networks, 5(2):157-166, 1994.
Cloud, G. Google cloud tpu v5e: Next- hardware for large-scale model https:://cloud..oogll/
generation ai training. bl blog/ http:p/potcttttaa-acccn--eeeeeee introducing-teoo-vai 2023. Accessed: 2024-09- 25.
Bisk, Y., Zellers, R., Le Bras, R., Gao, J., and Choi, Y. Piga: Reasoning about physical commonsense in natural language. In Proceedings of the Thirty-Fourth AAAI Conference on Artificial Intelligence, 2020.
Dahl, G. E., Schneider, F., Nado, Z., Agarwal, N., Sastry, C. S., Hennig, P., Medapati, S., Eschenhagen, R., Kasim- beg, P., Suo, D., et al. Benchmarking neural network training algorithms. arXiv preprint arXiv:2306.07119) 2023.
Blanchard, P., Higham, D. J., and Higham, N. J. Accurate computation of the log-sum-exp and softmax functions. arXiv preprint arXiv: / 034903469, 2019.
Boyd, S. and Vandenberghe, L. Conver optimization. Cam- bridge university press, 2004.
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory-effiient exact attention with io- awareness. Advances in Neural Information Processing Systems, 2022. URL https: /ttp/aaxvvvv///// org/ abs/ 2205.14135.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., and Maclaurin, D. JAX: Autograd and XLA. btttps://gite //gtttb..
jax, 2018. Ac-
cessed: 2024-09-25.
Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A. P., Caron, M., Geirhos, R., Alabdulmohsin, I., et al. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning, pp. 7480-7512. PMLR, 2023.
Bridle, J. S. Probabilistic interpretation of feedforward classification network outputs, with relationships to statis- tical pattern recognition. In Neurocomputtin:: Algorithms, architectures and applications, pp. 227-236. Springer, 1990.
Deng, J., Dong, W., Socher, R., Li, L.-.,, Li, K., and Fei-Fei, I. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255 255. Iece. 2009.
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J.. Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu,
Devlin, J., Chang, M.-W,, Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. arXiv preprint 1XXv:1810.04800 2018.
J.. Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M.. Gray, S., Chess, B., Clark, J., Berner, C., McCan- dlish, S., Radford, A., Sutskever, I., and Amodei, D.
10

---

Panayotov, V., Chen, G., Povey, D., and Khudanpur, S. Librispeech: An asr corpus based on public domain au- dio books. In 2015 IEEE International Conference on LASER: Attention with Exponential Transformation
Dosovitskiy, A., Beyer, L.,, Kolesnikov, A., Weissenborn,
Kim, H., Papamakarios, G., and Mnih, A. The lipschitz constant of self-attention. In International Conference on Machine Learning, pp. 5362-5571, PMLR, 2021.
D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. /929, 2021.
Kocijan, V., Chamorro- E., Sileo, D., Raiman, J., and Clark, P. Winogrande: An adversarial winograd schema challenge at scale. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pp. 8732-8740, 2020.
Dozat, T. Incorporating nesterov momentum into adam. 2016.
Gemini, T. Gemini 1.5:
Unlocking multimodal under- standing across millions of tokens of context. arXiv preprint arXiv:2403.05330 2024. URL https:// arxiv. org/ abs
Lai, G., Xie, Q., Liu, H., Yang, Y., and Hovy, E. Race: Large-scale reading comprehension dataset from exami- nations. In Proceedings of the 2017 Conference on Em- pirical Methods in Natural Language Processing, Pp. 785-794, 2017.
Gilmer, M., Dahl, G. E., Nado, Z., Kasimbeg, P., and Medapati, init2winit: a jax codebase for initialization, optimization, and tuning research, 2023. URL http: //github.com http:l/ccmmoooooooo.//.. http//ooolleiiinee Soogggeeinntiiinniii
LeCun, Y., Bottou, L., Orr, G. B., and Müller, K.-R. Effi- cient backprop. In Neural networks: Tricks of the trade, pp. 9-50. Springer, 2002.
Glorot, X. and Bengio, Y. Understanding the difficulty of training deep feedforward neural networks. In Pro- ceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256. JMLR
Levesque, J., Davis,
E.,, and Morgenstern, L.. The wino- grad schema challenge. In Proceedings of the Thirteenth International Conference Principles of Knowledge Representation and Reasoning, pp. 552-561,2011
Workshop and Conference Proceedings, 2010.
Google. Praxis: attentions.py, 2025. URL https:
Loshchilov, I. and Hutter, Sgdr: Stochastic gra- dient descent with warm restarts. arXiv preprint /008::/600.398.. 2016.
ccc0BeffeeSS34440eee222226666610 feexl///layers/attetiinno cessed: 2025-01-31.
Ac-
Loshchilov, and Hutter, Decoupled weight decay regu- larization. arXiv preprint arXiv:1711.05100 2017.
Google, R. Pax: A JAX-based neural network train- ing framework. https: httpgtttbbb///////////////t..htt cwm. google/ paxm1, 2023. Accessed: 2024-09-25.
Meta-AI, T. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024. URL sttps://arxiss 0ttg/abs/2407..22
Gulati, A., Qin, J., Chiu, C.-.,, Parmar, N., Zhang, Y., Yu, J. Han, W., Wang, S., Zhang, Z., Wu, Y., et al. Con- former: Convolution-augmenned transformer for speech recognition. In Proc. Interspeech, 2020.
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can suit of armor conduct electricity? new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 2381-2391, 2018.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual leamn ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
MLCommons. Bert dataset documentation. https: wttp:///wwwoooooooo//eooaiiiiiia
Hochreiter, and Schmidhuher, Long short-term memory. Neural computation, 98):1735-1780,, 1997.
masttrr/language_mdde///eee bert md. Accessed: 2024-10-16.
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, Transformers are RNNs: Fast autoregressive transform- with linear attention. In International conference machine learning, pp. 5156-5165. PMLR, 2020.
Mostafazadeh, N., Chambers, N., He, X., Parikh, D., Batra, D., Vanderwende, I.., Kohli, P., and Allen, J. A corpus and cloze evaluation for deeper understanding of com- monsense stories. In Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, pp. 839-849, 2016.
Khashabi, D., Chaturvedi, S., Roth, M., Upadhyay, S., and Roth, D. Looking beyond the surface: A challenge set for reading comprehension over multiple sentences. Pro- ceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 252-262, 2018

---

LASER: Attention with Exponential Transformation
Acoustics, Speech and Signal Processing (ICASSP), pp. 5206-5210. IEEE, 2015.
Zhang, S., Liu, H., Liu, S., Wang, Y., Liu, J., Gao, Z., Xu, W., Xu, Y., Sun, X., Cui, L.., al. Record: Bridging the gap between human and machine commonsense reading comprehension. arXiv preprint arXiv:/810.12885 2018.
Pilchvar, M. T. and Haimach-Coollaos,.. J. Wic: The word-in- context dataset for evaluating context-sensitive meaning representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, Volume (Long and Short Papers), pp. 1267-1273, 2019.
Radford, A., Narasimhan, K., Salimans, T., and Sutskever, Improving language understanding by generative pre- training. 2018. OpenAI.
Raffel, C.,, Shazeer, N., Roberts, A., Lee, K., Narang. S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. In Journal of Machine Learning Research, pp. 1-67, 2020. URL http: /mmrr./mr..///////////////////// org/ papers/v21/20- 074 .tmml.
volume 21, pp.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, and Polosukhin, I. Atten- tion is all you need. In Advances in neural information processing systems, volume 30, 2017.
Veličković, P., Perivolaropoulos, C., Barbero, F., and Pas- canu, R. softmax not enough (for sharp out-of- distribution). arXiv preprint arXiv:2410.01 104, 2024.
Wang, A., Pruksachatkun, Y., Nangia, N., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. R. Super- glue: A stickier benchmark for general-purpose language understanding systems. Advances in Neural Infor- mation Processing Systems (NeurIPS), pp. 3261-3275. Curran Associates, Inc., 2019.
Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Lan, Y., Wang, J., and Liu, T.-Y. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pp. 10524-10533. PMLR, 2020.
Ye. I., Dong, L., Xia, Y., Sun, Y., Zhu, Y., Huang, G., and Wei, F. Differential transformer. arXiv preprint arXiv:2410.05255 2024.
You, Y., Li, J., Reddi, S., Hseu, J., Kumar, S., Bhojana- palli, S., Song, X., Demmel, J., and Hsich, C.-J. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: 1904 904.00962, 2019.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Asso- ciation for Computational Linguistics, pp. 4791 4800, 2019.
12

---

LASER: Attention with Exponential Transformation
A. Appendix
A.I. Hyperparameter search space
In Table we outline the hyperparameter search space for all the benchmarks in Section 4.3.
Parameter
Min
Max
Scaling/Feasible Points
Learning.rate
10-4
10 2
log
1-81
10-2
0.15
log
8
0.9, 0.99, 0.999
warmup_factor
0.05
weight-decay
x
10
1.0
log
amboh.inot
0.1.
dropout_rate:
0.1
Table Hyperparameter search space used Section 4.3.
Parameter
Value
learning.rate
[le-1, le-2, le-3, le-4, le-5]
weight.decay
[le-2, le-1, 1.0]
beta_l
0.9
beta.2
epsilon
Ie-24
dropout.rate
0.0
Table 9. Hyperparameter search space language modeling experiments, Section 4.1
A.2. Proofs
Proof of Lemma 3.1. The softmax activation function is applied row-wise the preactivations Á; expand this computation row-wise follows:
softmax(A (A)
softmax(ā]] (āT)
softmax(ā.))
d softmax(ā:),
i (1,...,NN NJ
exp(ās1)
exp((r=)
En exp(āu)
En exp(ā.x)
exp(āii)
0ij = Σn exp((sx)
Taking gradient with respect to à; in the last expression gives:
j dai az (1 atz) if ifl= j
Dā,
Uijda else
Putting everything together, the Jacobian of the transformation a; softmax(ā:) can be written follows:
Daz
(diag(a:) - a;a,])
08
(diag(a;) a,a;)) Oe
(10)
0u;
Da,
13

---

LASER: Attention with Exponential Transformation
A.3. Training loss curves
3.0
3,0
LASER attention
LASER attention
Standard attention
Standard attention
2.8
2.8
Itain Loss
2.7
1 Loss
2.7
2.6
2.4
2.4
2.3
2.3
80000 Step
Step
(a) Train loss 2.2 billion parameter model.
(b) Test loss 2.2 billion parameter model.
2.80
LASER attention
2.75
LASER attention
2.8
Standard attention
Standard attention
2.70
2.65
Train Loss
Test Loss
2.60
2.0
2.55
2,50
2.5
2.45
244
2.40
20000
40000
00000
80000
100000
120000
140000
160000
2.35
20000
40000
60000
80000
100000 120000 140000 160000 Step
Step
(c) Train loss 1.1 billion parameter model.
(d) Test loss 1.1 billion parameter model.
Figure 5. Performance comparison for 2.2 billion and 1.1 billion parameter models. The 2.2 billion model has 32 layers, 8 attention heads with head size 512, MLP hidden dimension 8192, and model dimension 2048. The 1.1 billion model layers, attention heads (head size 256), MLP hidden dimension 4096, and model dimension 1024. LASER demonstrates better train and test loss compared to Standard Attention even in large-scale settings.
A.4. Fluctuations in downstream evaluations
(a) Part
(b) Part 2
Dataset
LASER
Standard
Dataset
LASER
Standard
WSC
80.98
79.30#0.44
RaceH
37.9910.13
37.5610.11
Winogrande
0.21
62.15
PIQA
77.20
76.74.0.12
Winograd
0.37
80.2910.27
OpenBookQA
49.0810.20
17.56.00155
WiC
0.66
51.29+0.55
MultiRC
57.16110.18
53.9910.116
StoryCloze
0.09
76.46+0.07
HellaSwag
66.61+0007
65.46
RTE
53.14410.14
53.1410.14
COPA
82.0010.00
80.00
ReCoRD
8.02
85.06+0.05
CB
40.00++0.87
43.93+0.87
RaceM
0..5
49.6740.13
BoolQ
63.3910..8
60.45
Table 10. Here add standard deviations of each number, computed downstream evaluations of checkpoints and note average for LASER to be 63.58.10.26 and an average for standard attention to be 60.69.10.23, noting a difference 0.89.
14

---

LASER: Attention with Exponential Transformation
A.5. Training Analysis
Training instability. There can be spikes in training loss curves initially during large language model training. We notice that despite these spikes training stabilizes and converges smoothly. However, training instability/spikes be attributed to poor model architecture and optimizer choices. We now ablate the choice of attention mechanism and understand its affect on training stability. Figure 6 compares training stability of different models.
Lecen attoninn
LASER abenten Strecert atterton
Dxpp
2tp
Stp
(a) 234M Parameters
(b) 300M Parameters
(c) I.IB Parameters
(d) 2.2B Parameters
Figure Train loss vs. steps LASER and standard attention across different parameter scales. Training stability for each attention mechanism can be observed through the number of training spikes. Generally, models with LASER attention exhibit fewer training spikes compared to models with standard attention, indicating greater stability in training for LASER attention across all parameter seales. The figures focus the initial part of the training as the remainder did not exhibit training instability.
Additionally, we ran pretraining of a 2.2B parameter language model presented in Section 4.1 with BF16 precision weights and activations and compare LASER and Standard attention in Table 11.
Precision
LASER Loss
Standard Attention Loss
FP32
2.326
2.344
BF16
2.333
2.350
Table 11. Comparison of LASER and Standard Attention Losses for FP32 and BF16 precision 2.2B model.
We cached 4800 query, key, value matrices of size (1024, 8, 256) (headsize-2555 256, heads-8, sequence length- 1024) during training, and computed the following numerical reconstruction relative errors of attention output in b vanilla LASER (0.0018, 0.0002), LASER+log- wwweehhed-suum-cpppeeeeee experiment conducted in machine with 4 TPUv5 chips. We conducted the same experiment on 16 A100s and found the following errors: vanilla LASER (0.002, 0.0003), (0.0018, 0.0002). While on an average, found log- weiighed--sum-ex--teee that this trick prevents overflows, which is crucial for stable training. Similar trick famously known as log-sum-exp is used to prevent overflows due to exp(.) function in the softmax function and is adopted by both Pytorch and Jax in their softmax implementations.
(0.0017,0.0001)) Standard Attention (0.0016, 0,0001). This
(0.0019,0.000 1,000,,, Standard Attention to help on both TPUv5 and A A100s,
note
Understanding training instability. The enhanced training stability of the LASER attention mechanism can be formally attributed its direct solution for the gradient saturation problem inherent in standard attention models. The Jacobian of the softmax function, which scales the gradients for the query and key projection matrices (Wq and Wx). diminishes towards zero when attention probabilities become sharp. This vanishing gradient signal create inconsistent updates to parameters, leading to suboptimal learning and potential loss spikes. LASER circumvents this by reformulating the attention output as log(softmax (QKT) exp(V)) a structure which admits with lower saturation properties.
This improved parameter learning grants LASER its stability, a behavior analogous to residual connections in deep networks. The LASER output for given token i, oz can be expressed as a log-sum-exp function, which tight, differentiable approximation of the max function. Here we bound 01:
max(1; [1 + log(an)) 02 + log((12))
< 01 I max(11 + log(an1), 22 + log(a12)) + log(2).
In cases of strong self-attention, where the attention probability az is high, the max function property ensures that the output o; closely approximates its corresponding input value vy. By creating this effective "pass-through" for the input
15

---

LASER: Attention with Exponential Transformation
signal during moments of high certainty, LASER functions as an inherent residual connection. We conjecture that this property is the formal mechanism that underpins the empirically observed reduction in training loss spikes compared to standard attention.
16


Reviews:
[{"review_id": "6QDcRnntsQ", "reviewer": "Reviewer_tFgw", "summary": "This paper introduces LASER (LogArithm of Summed Exponentials of Representations), a novel attention mechanism for Transformers. The researchers found that in the standard attention mechanism, the gradients backpropagated through the softmax operation can be small, which may lead to inefficient learning of parameters. LASER addresses this by applying attention in the exponential value space. It conducts attention on exp(V) and uses a Log - Weighted - Sum - Exp trick to avoid numerical overflow. Experiments on various Transformer models, including autoregressive LLMs, Vision Transformer, and Conformer, show that LASER can improve performance. For example, in autoregressive language modeling on the C4 dataset, it can achieve up to a 1.74% relative improvement in test loss, and it also shows better results in downstream tasks with an average accuracy improvement of about 1%.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "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": "We thank the reviewer for their comprehensive comments.\n\n>  The paper mainly focuses on models with a certain scale of parameters, and it's not clear how well it will perform in models with extremely small or large numbers of parameters. \n\nWe trained a model with **7.7B parameters scaling up by 3.5x our largest model (2.2B)**. We train the model for 44B tokens and evaluate it on downstream tasks:\n\n| Dataset     | Laser  Acc | Standard  Acc |\n|-------------|----------------|--------------------|\n| arc_e       | 52.48          | 52.69              |\n| boolq       | 62.45          | 56.48              |\n| cb          | 44.64          | 42.86              |\n| hellaswag   | 57.16          | 56.02              |\n| multirc     | 56.00          | 55.59              |\n| openbookqa  | 45.40          | 43.60              |\n| racem       | 44.64          | 44.15              |\n| rte         | 53.43          | 49.82              |\n| storycloze  | 71.78          | 71.51              |\n| wic         | 47.34          | 47.34              |\n| winogrande  | 58.41          | 57.77              |\n| **average** | **53.97**      | **52.53**          |\n\nWe found that compared to 2.2B billion model which gave an gain of 1% on average, for 7.7B model, we found about **1.44% average gain over standard attention**, and observe **more pronounced differences in boolq (+6%), cb (+1.78%), openbookqa (+1.8%), rte (+3.61%)**.  In 7.7B parameter model, we scaled along all dimensions, model_dimension=3440, hidden_dimension=11584, num_heads=16, dims_per_head=720. \n\nWe conducted a power-law fit (Figure 4) on test loss of models ranging from 234M to 2.2B and conjecture that **smaller models might have smaller differences in loss value with LASER**. However, this observation is **for decoder language models on C4**. In contrast, our experiments with **ViT-S/16 with 22M parameters, LASER gives 1.2% improvement** in classification error. \n\n> Also, it might face challenges in some specific tasks where the data characteristics are very different from those in the experiments.\n\nWhile we showed gains on standard downstream tasks used to evaluate the performance large language models, identifying tasks where LASER performs poorly and similarly softmax-attention performs poorly is interesting future work we plan to do. \n\n\n> The Log - Weighted - Sum - Exp trick is crucial for LASER to avoid numerical overflow. But this also means that the performance of LASER is highly dependent on this trick.\n\n\n**We cached 4800 query, key, value matrices of size (1024, 8, 256) (headsize-256, heads-8, sequence length-1024) during training**, and computed the following numerical reconstruction relative errors of attention output in bfloat16: **vanilla LASER (0.0018, 0.0002),  LASER+log-weighted-sum-exp-trick (0.0017,0.0001), Standard Attention (0.0016, 0.0001)**. This experiment is conducted in a machine with 4 **TPUv5 chips.** We conducted the same experiment on **16 A100s** and found the following errors:  **vanilla LASER (0.002, 0.0003),  LASER+log-weighted-sum-exp-trick (0.0019,0.0002), Standard Attention (0.0018, 0.0002)**. While on an average, we found log-weighted-sum-exp-trick to help on both TPUv5 and A100s, **we note that this trick prevents overflows, which is crucial for stable training**. Similar trick famously known as **log-sum-exp [1] is used to prevent overflows due to exp(.) function in the softmax function** and is adopted by both pytorch and jax in their softmax implementations.\n\n>  There are many advanced attention mechanisms proposed recently, and the paper doesn't compare LASER with them.\n\n**Diff Transformers [1] is a recent state of the art attention and concurrent work, which outperforms Standard Attention** which uses the difference of two softmax attention matrices and their multiplication with the value matrix as the output. We note that both matrices still face backpropagation challenges due to the use of softmax. To address this, we use the LASER formulation. We trained a 2.2B DiffTransformer with Model dim: 2048, Hidden dim: 8192, Number of attention heads: 8, Head size: 512, Number of layers: 32. The model was trained on 24 billion tokens. \n| Dataset     | Diff+LASER | DiffTransformer |\n|-------------|-------------|-----------------|\n| arc_e       | 49.2845     | 49.2003         |\n| cb          | 42.8571     | 41.0714         |\n| hellaswag   | 51.9319     | 51.5834         |\n| multirc     | 55.1361     | 52.8259         |\n| openbookqa  | 44.2000     | 43.0000         |\n| racem       | 42.4095     | 40.7382         |\n| rte         | 52.3466     | 50.9025         |\n| storycloze  | 71.0315     | 71.2988         |\n| wic         | 50.0000     | 49.5298         |\n| winogrande  | 56.0379     | 55.8011         |\n| **Average** | **51.5235** | **50.5951**     |\nOn average, **we observe an improvement of ~1%, similar to the improvement in 2.2B transformer model in Section 4.1**.\n\n[1] Ye, Tianzhu, et al. \"Differential transformer.\"  (2024).", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "no", "essential_references_not_discussed": "no", "other_strengths_and_weaknesses": "Strengths\n1. LASER really shows its worth in experiments. It can significantly reduce the test loss in autoregressive language modeling, like up to 1.74% in some cases. And it also improves the accuracy in downstream tasks, which means it can make the model perform better in practical applications. For instance, in the Vision Transformer on Imagenet, it has a 1.15% absolute improvement in error rate, which is a great result.\n2. It's really convenient that LASER can be implemented with just small modifications to the existing attention implementations. It doesn't need to change the underlying attention function, which makes it easy for researchers and engineers to apply in their work. This simplicity also means it can be quickly integrated into different models.\n3. The paper provides a solid theoretical analysis. By studying the gradient backpropagation in the attention mechanism, it clearly points out the problem of small gradients in the standard attention and explains how LASER can solve this problem. This theoretical support makes the proposed method more reliable and convincing.\n4. LASER has been tested on a variety of models across different modalities, such as text, speech, and vision. It shows consistent improvements in all these models, which indicates that it is a versatile method and not limited to a specific type of model or task. This broad applicability makes it more valuable in different research and application scenarios.\n\nWeaknesses\n1. Although LASER shows good results in many experiments, it may not be suitable for all situations. The paper mainly focuses on models with a certain scale of parameters, and it's not clear how well it will perform in models with extremely small or large numbers of parameters. Also, it might face challenges in some specific tasks where the data characteristics are very different from those in the experiments.\n2. The Log - Weighted - Sum - Exp trick is crucial for LASER to avoid numerical overflow. But this also means that the performance of LASER is highly dependent on this trick. If there are some issues with this trick in different hardware or software environments, it may affect the performance of LASER. And the paper doesn't fully explore the potential problems that might occur when using this trick.\n3. The paper mainly compares LASER with the standard attention mechanism and some simple modifications of it. There are many advanced attention mechanisms proposed recently, and the paper doesn't compare LASER with them. So, it's hard to say how LASER stands out among the most advanced techniques in the field. This lack of comparison limits the understanding of LASER's superiority and competitiveness.", "other_comments_or_suggestions": "no", "questions_for_authors": "see weaknesses", "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": "", "extra_scores": {}}, {"review_id": "AXSidrwXbM", "reviewer": "Reviewer_8gDY", "summary": "The papers studies the gradients during backpropagation of the standard softmax dot product attention within transformers. The authors' key insight is that these gradients can be extremely small, leading to poor gradient signal propagation, which in turn leads to sub-optimal learning. They then suggest a modification of attention to mitigate this issue by suggesting that attention should be carried into in an exponential value space, involving applying an exponential function to the values matrix along with a log function to the softmax probabilites.\n\n\n**Update After rebuttal:**\n\nI thank the authors for their rebuttal. I have re-read their paper several times including the appendix. I have also read the other reviews along with their rebuttals to the other reviewers and to my review. However, I am still not convinced that this work has enough novelty and significance to be published in a conference at the level of ICML.  Therefore, I will be voting to reject the paper.\n\nA real concern I have with the paper is that they have shown a lot of empirical results yet offer no real theoretical understanding of why their method works in the way they have reported. Furthermore, the paper has so many statement that are well known to researchers that publish in this area. For example, the whole of section 3.1 and 3.2 is well known and yet they dedicate more than a page to it. Lemma 3.1 is a very simple lemma which they claim they extend to arbitrary sequence length $N$ yet the extension is extremely simple this cannot be considered a contribution. Furthermore, the theory of the paper simply consists of computing gradients and showing that certain quantities from the computation yield a saturation term (possibly of low order). This is by no means novel or significant.\n\n**Final Comments on Rebuttal:**\n\n1. **We would like to highlight a subtle distinction in our approach: we apply the logarithm after the multiplication..:**\n\nThank you for clearing that up. I am still not convinced that this is novel. Just because an idea has not been explored that does not make it novel.\n\n2. **Please refer to our response to Reviewer tFgw, where we scaled to 7.7B parameter models and found an average improvement of 1.44%...:**\n\nYou empirical results are fine but there is a major issue in that there is no real explanation in why you are able to get these improvements.\n\n3. **With better gradient signal propagation, which we theoretically demonstrate in our manuscript, we conjecture that gradients may be better conditioned with LASER than with Standard attention. However, we have not yet explored the effect of this on the convergence rate of the optimizer.**\n\nI do not agree with this entirely. You actually do not prove that your method gives better signal propagation. You simply compute gradients and show how you can then apply your method to avoid the saturation that can take place. There is no proof that what you do leads to better signal propagation and thus the relationship to the empirical results is thus still lacking. You should really prove that your method leads to better conditioned gradients. That will then clearly show your method is capping off any over flow and leading to better training.\n\n4. **We observed similar spikes in gradient norms for both LASER and Standard attention, but only a few such spikes translated into actual loss spikes. We address this in our stability analysis (Appendix A.5), where we note that Standard attention suffers from more loss spikes than LASER.**\n\nThank you. I noticed massive spikes in your analysis (appendix A.5) even with normal attention which you have not made clear as to why that is happening. The fix could simply be that a better learning rate scheduler need to be chosen. Yet no analysis is done to form an actual conclusion.\n\n5. **Thank you for pointing this out. We will explicitly mention the following limitations of LASER in our revision...**\n\nThe quadratic complexity holds even for normal attention so really you have not added anything by saying that point. I was asking about explicit limitations of LASER that are not present in usual attention. The fact that there were no limitations stated and analysed in the original paper is quite concerning, especially when reading the paper I noticed several limitations as stated above.", "questions": "", "limitations": "", "rating": 1, "confidence": null, "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": "We appreciate the reviewer’s insightful feedback.\n\n> **Reviewer:** The authors simply show that gradients saturate in the attention gradients during backpropagation and then show that by applying an exponential on the values and a logarithm on the softmax probabilities one can produce a gradient that does not saturate.\n\nWe would like to highlight a subtle distinction in our approach: we apply the **logarithm *after*** the multiplication of softmax probabilities with the values. To the best of our knowledge, this idea has **not been explored** in existing literature, and we believe it offers a **novel attention formulation**.\n\n---\n\n> **Reviewer:** Significance: This is related to the novelty. I don't believe the paper will be significant.\n\nPlease refer to our response to **Reviewer tFgw**, where we scaled to **7.7B parameter models** and found an **average improvement of 1.44%** in downstream evaluations, with **more pronounced improvements up to 6%** in individual tasks. Additionally, we show that this insight **generalizes to other tasks** such as:\n- **ViT: +1.2%**\n- **BERT**\n- **Speech Transformer**\n\nAlso, in our response to **Reviewer 9sxs**, we fine-tuned the **2.2B model on the SuperGLUE dataset** [1] for 10 epochs and observed a **1.65% improvement in decoding accuracy** (in contrast to the eval/ranking accuracy reported in Table 2) across the SuperGLUE benchmarks.\n\n---\n\n> **Reviewer:** Would the authors be able to say something about whether theoretically LASER leads to a better convergence rate for the optimizer?\n\nWith better gradient signal propagation — which we theoretically demonstrate in our manuscript — we **conjecture** that gradients may be **better conditioned** with LASER than with Standard attention. However, we have **not yet explored** the effect of this on the convergence rate of the optimizer. We will consider including this analysis in our final revision.\n\n---\n\n> **Reviewer:** In Figure 3 I notice an issue with the gradient norm. Between step 60000 and 80000 I see that LASER has a sudden jump in gradient magnitude. This is usually an effect of gradient becoming too large which is not good for training. Did this happen in other experiments? This could be a limitation of the work.\n\nWe observed similar spikes in **gradient norms for both LASER and Standard attention**, but only **a few such spikes translated into actual loss spikes**. We address this in our **stability analysis (Appendix A.5)**, where we note that **Standard attention suffers from more loss spikes than LASER**.\n\n---\n\n> **Reviewer:** What are the limitations of LASER? I noticed you did not really discuss any limitations of your method.\n\nThank you for pointing this out. We will explicitly mention the following **limitations of LASER** in our revision:\n- LASER attention has **quadratic computational complexity**, which can be significant for long sequence lengths.\n- A **comprehensive failure case study** for LASER (and similarly for Standard attention) has not been conducted. This remains **future work**.\n\n---\n\n[1] Wang, Alex, et al. *\"SuperGLUE: A stickier benchmark for general-purpose language understanding systems.\"* Advances in Neural Information Processing Systems 32 (2019).", "claims_and_evidence": "Yes the claims are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The methods and evaluation criteria are fine.", "theoretical_claims": "The theoretical claims both in the paper and the appendix were checked and as best as I can see are correct.", "experimental_designs_or_analyses": "The experimental designs are fine.", "relation_to_broader_scientific_literature": "The key contributions in the paper are related to transformer training and the well known issues of gradient propagation in neural networks. this is a well studied and important topic and the authors give a new viewpoint of this.", "essential_references_not_discussed": "References are fine.", "other_strengths_and_weaknesses": "**clarity:** I admire the authors efforts in trying to understand gradient propagation within the attention layer of a transformer. Their approach of mapping everything to exponential space and working there is a nice, although somewhat simple, insight. The clarity with which they write the paper is simple to read.\n\nHowever, I do have the following issues:\n\n**Novelty:** I don't feel there is enough novelty in this paper to be published in ICML. The authors simply show that gradients saturate in the attention gradients during backpropagation and then show that by applying an exponential on the values and a logarithm on the softmax probabilities one can produce a gradient that does not saturate. The approach is simply to look at the saturated gradient part and realise that a certain transformation can mitigate the effect of the saturation. Furthermore, the mathematical approach is simply just computing a derivative using the chain rule. What would have been much better is if the authors actually showed that this saturation caused serious problems with convergence of the optimizer. Or for that matter showing that by getting rid of saturation by using their LASER method that a different optimizer such as SGD can be used to train a transformer as it is well known that SGD performs extremely badly when compared to Adam on many transformers. My feeling is this will lead to nothing but the authors are welcome to prove me wrong.\n\n**Significance:** This is related to the novelty. I don't believe the paper will be significant. While their insight is nice. It is my feeling that that is all it is. A nice insight that leads to some better performance with a simple adjustment done by carrying out some simple chain rule calculations. I feel the authors should have really shown how LASER can be impactful and significant by showing that it leads to something practitioners in the field did not expect with transformers such as better training with different optimizers like SGD.", "other_comments_or_suggestions": "I do not have any other comments.", "questions_for_authors": "1. Would the authors be able to say something about whether theoretically LASER leads to a better convergence rate for the optimizer?\n2. Furthermore, I don't really see a big difference between LASER and standard attention in the training curves. For example, in figure 2 (top curve) it seems LASER converges faster but only slightly. Since your work is that LASER should have better back propagated gradients shouldn't this lead to a much better and faster train loss curve?\n3. In figure 3 I notice an issue with the gradient norm. Between step 60000 and 80000 I see that LASER has a sudden jump in gradient magnitude. This is usually an effect of gradient becoming too large which is not good for training. Did this happen in other experiments? This could be a limitation of the work.\n4. What are the limitations of LASER? I noticed you did not really discuss any limitations of your method.", "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": "", "extra_scores": {}}, {"review_id": "Bijku08WhH", "reviewer": "Reviewer_9sxs", "summary": "This paper studies the problem of Gradient saturation in softmax in the attention architecture. Thus, it proposes LASER, a log-sum-exp structure to replace the original dot-product formulation in attention with Log-Weighted-Sum-Exp Trick. Extensive experiments on different benchmarks and models, with in-depth analysis, support the development of such a novel attention mechanism.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "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": "Thank you for the reviewer’s helpful observations.\n\n> I'm curious about why LASER is stabler than vanilla attention mechanism at training time, since the author mentioned that LASER introduces larger gradient norm, intuitively this may cause fluctuation in training. Can the author provide some analyses on this from perspectives like gradient distribution?\n\nTo understand this, we draw analogy to a residual connection. While residual networks add residual connections to improve gradient backpropagation, residual connections increase the gradient norm compared to vanilla network, however, the stability improves as more parameters are activated. Analogously, we conjecture that, fixing the gradient backpropagation through query and key within softmax improves the learning of W_K and W_Q. \n\n\n> I'm open to discussion for the above-mentioned issues, and I think the justification of LASER in terms of some analyses on the better stability could make the paper more thorough.\n\nWe will formally prove the above in our final revision, using the aforementioned intuition.\n\n> The quantitative experimental results on downstream tasks demonstrate marginal improvements compared to the vanilla attention mechanism.\n\nIn our response to Reviewer tFgw, we trained a 7.7B model and observed **average improvements of 1.44%**, with more pronounced gains of up to **6%** on individual downstream tasks. \n\nAdditionally, we fine-tuned the 2.2B model on the SuperGLUE dataset [1] for 10 epochs and found a **1.65% improvement in decoding accuracy** (as opposed to eval/ranking accuracy reported in Table 2) on the SuperGLUE benchmarks listed below:\n\n| Task  | LASER | Standard |\n|-------|--------|----------|\n| Copa  | 57.00  | 58.00    |\n| Wic   | 56.64  | 53.92    |\n| WSC   | 40.38  | 36.54    |\n| RTE   | 22.02  | 20.94    |\n| **Avg** | **44.01** | **42.35** |\n\nWe believe these new fine-tuning results offer a more comprehensive view of how LASER improves over standard attention, beyond what is seen in 1-shot downstream evaluations.\n\n\n\n[1] Wang, Alex, et al. *\"SuperGLUE: A stickier benchmark for general-purpose language understanding systems.\"* Advances in Neural Information Processing Systems 32 (2019).", "claims_and_evidence": "All the claims are correct with experimental / theoretical analysis as evidence.", "methods_and_evaluation_criteria": "The method is compared with various baseline choices, adopted to different models on different tasks and is universally better than baselines on average. However, one possible concern is the improvements are marginal. Can the author justify the use of LASER in terms of other perspectives, e.g. faster convergence during training, stabilizing training, etc. While the author provides some results on the training stability in the Supplementary materials, it is appreciated if the author can provide some analysis on why LASER is stabler than original attention mechanism.", "theoretical_claims": "All the theoretical claims appear to be correct.", "experimental_designs_or_analyses": "The experimental designs and analyses are thorough, especially for the part where different optimizers are ablated.", "relation_to_broader_scientific_literature": "The key contribution of this paper, a new attention mechanism that aims to solve the gradient saturation problem in attention, has great potential for a great impact in attention-based models, which are dominant in the current era.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n\n1. The proposed LASER method indeed theoretically solves the gradient saturation problem.\n2. The loss curves showed empirically show the effectiveness of LASER.\n\nWeaknesses:\n1. The quantitative experimental results on downstream tasks demonstrate marginal improvements compared to vanilla attention mechanism.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "I'm open to discussion for the above-mentioned issues, and I think the justification of LASER in terms of some analyses on the better stability could make the paper more thorough.", "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": "", "extra_scores": {}}, {"review_id": "XGneURtbAV", "reviewer": "Reviewer_o5d3", "summary": "This paper addresses the problem of vanishing gradients in standard dot-product attention in transformers. The authors show mathematically how this problem arises in the Jacobian of the attention function during backpropagation. Next, a new technique called LASER (Logarithm of Summed Exponentials of Representations) is introduced. As the name suggests, this technique, at it's core, is like the log-sum-exp function which is a new way of formulating/approximating the weighted average in attn(.) using the log of summed exponentials. The benefit of this technique is the avoidance of the vanishing gradients as shown mathematically in this paper and consequently the better training of transformer based models.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "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": "We appreciate the reviewer’s comments and suggestions.\n\n\n### Addressing Other Comments or Suggestions\n\n**Comment 2:**  \nUnfortunately, this is a typo — it is supposed to be `exp(v_1) - exp(v_2) >> 0`. While `v_1 - v_2` might not be significantly different, the difference between `exp(v_1)` and `exp(v_2)` can be much larger due to the exponential function, as dictated by the value parameter \\( W_V \\).\n\n**Comments 1 & 3:**  \nThank you for pointing these out. We will correct them in our revision.\n\n**Comment 4:**  \nThanks for pointing out this typo — there should be no `exp(.)` in that expression. We appreciate the feedback and will fix this in the revision.\n\n**Comment 5:**\nWe conjecture that, since half the attention scores are zeros in decoder-only models, this could lead to less dependence on W_K and W_Q parameters, and less gradient flow through W_K and W_Q compared to BERT models.  \n\n---\n\n### Questions for Authors\n\n> **Reviewer:** My main concern about the paper is: how significant are the improvements shown? In all the different tasks, the improvements seem minor. Did the authors perform statistical significance testing?\n\nIn Table 7 (Appendix A.4), we note that the **standard deviation is low** compared to the difference between LASER and Standard attention, indicating that the observed improvements are statistically significant.\n\n> **Reviewer:** Furthermore, how does LASER scale to even larger transformer models? I understand that it might not be possible to conduct that experiment but the authors should comment on how LASER can be useful for larger models given how minor the improvements are.\n\nPlease refer to our response to **Reviewer tFgw**, where we trained a **7.7B model** and demonstrated **average improvements of 1.44%**, with more pronounced gains of up to **6%** on individual downstream tasks.\n\nAdditionally, please see our response to **Reviewer 9sxs**, where we fine-tuned the **2.2B model** on the **SuperGLUE dataset** [1] for 10 epochs and found a **1.65% improvement in decoding accuracy** (as opposed to eval/ranking accuracy in Table 2).\n\n> **Reviewer:** Related to the above, does LASER introduce any computational overhead compared to vanilla attention? If yes, how is it worth using LASER given the performance and computational overhead tradeoff?\n\nAs noted in **Lines 326–329** under *Performance Analysis*:\n- The **2.2B model with standard attention** takes **27.22 hours** on TPU v5  to reach a minimum test loss of 2.327.\n- In contrast, **LASER takes only 24.88 hours**, a **relative walltime improvement of 9.4%**.\n\n---\n\n[1] Wang, Alex, et al. *\"SuperGLUE: A stickier benchmark for general-purpose language understanding systems.\"* Advances in Neural Information Processing Systems 32 (2019).", "claims_and_evidence": "Yes, claims are clear with convincing evidence.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I did not check the proofs.", "experimental_designs_or_analyses": "Yes. The experimental designs are sound and valid.", "relation_to_broader_scientific_literature": "This paper is a key contribution to the broad literature of attention mechanisms in transformers.", "essential_references_not_discussed": "The paper does a good job with citing the essential references.", "other_strengths_and_weaknesses": "Strengths:\n\nThis paper address a key issue in transformer models with a simple yet novel and innovative idea. The paper is very well written and the flow of ideas are easy to follow. The paper is also mathematically grounded with experiments spanning across multiple applications.\n\nWeakness:\n\nThe performance improvements shown are minor.", "other_comments_or_suggestions": "1. In section 3.2, the definition of $\\tilde{A} = KQ^T$ is different from definition in 3.1\n2. Is $v_1 - v_2 >> 0$ an obvious choice? Is it always true?\n3. Equation (8) should be $log(exp(v_1 + log(a_{11})\\textbf{)}+exp(v_2+log(a_{12}))\\textbf{)}$ (the text is missing two closing brackets)\n4. In algorithm 1, step 4: what is $exp(.)$ ?\n5. In section 4.2, authors say that LASER shows more improvement in autoregressive LM compared to BERT. Is there some explanation/hypothesis to this in light of how LASER is formulated?", "questions_for_authors": "My main concern about the paper is: how significant are the improvements shown? In all the different tasks, the improvements seem minor. Did the authors perform statistical significance testing? \n\nFurthermore, how does LASER scale to even larger transformer models? I understand that it might not be possible to conduct that experiment but the authors should comment on how LASER can be useful for larger models given how minor the improvements are.\n\nRelated to the above, does LASER introduce any computational overhead compared to vanilla attention? If yes, how is it worth using LASER given the performance and computational overhead tradeoff?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: I believe the paper is worth accepting. While LASER's innovation in addressing gradient backpropagation through attention may be modest, the empirical evidence is compelling. However, I agree that wider field adoption would require more extensive practical validation beyond what's shown in the paper. The authors have responded well to reviewer concerns with additional experiments on larger models and diverse tasks. Those are a major factor that I consider when I recommend the acceptance.
Decision: Accept (poster)