Unifying Linear-Time Attention via Latent Probabilistic Modelling
Authors: Rares Dolga
Lucas Maystre
Marius Cobzarenco
David Barber
Venue: Accepted by TMLR ()

Abstract: Transformers have achieved state-of-the-art results across a range of domains, but their quadratic attention mechanism poses significant challenges for long-sequence modelling. Recent efforts to design linear-time attention mechanisms have yielded more scalable alternatives, yet often at the cost of performance, particularly on discrete data such as language. In this work, we revisit linear attention through the lens of probabilistic graphical models. We first show that standard linear attention can be interpreted as an undirected latent variable model, revealing a key limitation: the absence of directionality. To address this, we propose a novel directed parameterisation of linear attention that introduces an asymmetric structure, enabling an interpretation aligned with the causal and sequential nature of language. Our formulation integrates global latent-variable attention with local standard attention in a fully probabilistic framework. Additionally, we introduce a recurrent parameterisation of queries and keys that avoids reliance on relative positional encodings, often incompatible with linear attention. Experiments on language modelling benchmarks demonstrate that our model achieves competitive performance with standard attention and outperforms existing linear attention variants.

Published in Transactions on Machine Learning Research (12/2025)
Unifying Linear-Time Attention via Latent Probabilistic Modelling
Rares Dolga
rores. 1 www.cl..c.ccc...oo
University College London, AT Centre
UiPath
Lucas Maystre
hucas.mayysfn www.atrrrwuppphcccccc com
UiPath
Marius Cobzarenco
morius marius..cobarrenconnnipppttt
UiPath
David Barber
david. burber Buippthh.com com
University College London, AI Centre
UtPath
Reviewed on OpenReview: https: // owtp://veeprr.eee.e.mmmm net/ forum? id= 1R71
Abstract
Transformers have achieved state-of-the-art results across a range of domaīns, but their quadratic attention mechanism poses significant challenges for long-soquence modelling. Ro- cent efforts design linear-time altention mechanisms have yielded more scalable alterna- tives, yet often at the cost of performance, particularly discrete data such language. In this work, we revisit linear attention through the lens of probabilistic graphical models. We first show that standard linear attention can be interpreted as an undirected latent variable model, revealing a key limitation: the absence of directionality. To address this, we propose a novel directed parameterisation of linear attention that introduces an asymmetric struc- ture, cnabling an interpretation aligned with the causal and scquential nature of language. Our formulation integrates global latenl-variable atlention with local standard attention in fully probabilistic framework. Additionally, queries and keys that avoids reliance relative positional encodings, often incompatible with linear attention. Experiments on language modelling benchmarks demonstrate that our model achieves competitive performance with standard attent ion and outperforms existing linear attention variants.
introduce recurrent parameterisation of
1 Introduction
Transformers have become a cornerstone of modern deep learning, with attention mechanisms at the heart of their success. Despite their strong performance across wide range of tasks, the quadratic time complexity of standard attention imposes efficiency challenges when scaled extremely long sequences. This limitation has motivated extensive research into linear-time allention mechanisms. While several linear variants have been proposed (Katharopoulos et al.) 2020 Lin al. 2021), most fall short of the performance achieved by standard Transformers, particularly in language modelling.
Recent work highlights the importance of gating mechanisms for improving linear attention performance on discrete data such as text (Yang et al. 2024h. However, these approaches still build on the classical formulation of lincar attention, which approximates softmax attention by expressing it as an inner product of kernel feature maps and leveraging associativity of matrix multiplication. Although the standard attention mechanism be viewed probability distribution, the linear formulation lacks probabilistic foundation,

---

Published in Transactions on Machine Learning Research (12/2025)
and it remains open question whether stronger linear mechanisms can be developed by drawing probabilistic principles.
In parallel, many hybrid models combine linear global attention with local standard attention (Hua et al, 2022j. However, these methods often fail to properly normalize attention weights across the sequence, which can introduce biases, such as putting more weight on the local attention.
In this work, we revisit the foundations of linear attention through the lens of probabilistic graphical models. We first show that standard linear attention can be interpreted Building on this insight, propose a novel directed parameterisation of linear attention. We will show that although both the undirected and directed models encode the same conditional independence structure, the directed variant enables a generalive interpretation, where a token influences another token 8 via a latent variable l. This asymmetry aligns with the inherently directional nature of attention and is particularly well-suited for modelling sequential or causal relationships.
undirected latent variable model.
Conceptually, our approach can be viewed as clustering lokens around high-level latent concepts, with attention computed via interactions between tokens and these latent variable states. Our framework naturally supports the integration of global and local attention while maintaining properly normalized probabilistic structure.
Our main contributions are:
We provide probabilistic interpretation of standard linear attention undirected latent variable model.
Leveraging this insight, we introduce a novel directed parameterisation of linear attention that enables asymmetric modelling.
We show how to integrate standard local attention with our approach while preserving full normal- ization of attention weights.
Finally, propose recurrent parameterisation of queries and keys, replacing relative positional encodings (which are incompatible with linear allention), and demonstrate strong performance on language modelling dasks-comparable to standard attention and recent saate-of-the-are linear models.
2
Background
Viewing attention through a probabilistic lens provides a unified account of both standard and linear atten- tion and clarifies their relationships. Since attention is probability distribution, standard aitent-variable parameterisations are natural. This framing places common variants within approximate probabilistic in- ference, makes their implicit approximations explicit, and enables the use of classical methods grounded in latent variables.
From this perspective, we propose a directed (unsymmetric) parameterisation of linear attention that well suited to language. We further develop fully normalised hybrid that combines this directed variant with sliding-window standard attention and achieves strong empirical performance. These results demonstrate how the probabilistic view yields concrete methodological improvements.
2.1 Probabilistic Model for Attention
The attention mechanism takes an input sequence of token vectors X veat (xI...., æt); as E RD and transforms this a new sequence according to an input-dependent linear function:
X = Attention (X) = softmax (QK 8 M) V
where Q = XW, K= XW, V = XW., M is the causal attention mask or matrix of ones for bidirectional attention. WG, WE, W. DDD'
2

---

Published in Transactions on Machine Learning Research (12/2025)
We can now rewrite the vectorised Equation for each token in the sequence. Although we focus the causal case, the bidirectional case is similar.
Et
aes
(2)
==1
where Ots are the attention weights defined for the causal case as:
exp (q: ki.)
if
Ote E's exp (a: ki))
(3)
0
otherwise
Since Ots is a positive normalised quantity, we can interpret ats the probability p(s]t) of the token (i.e the sequence element) occurring at position given the token occurring at position t. While attention represents learning p(s]t) directly and using this quantity, the joint distribution p(s,t) = p(s]t)p(t) can be represented by the Markov Network in Figure
Figure 1: Graphical model for bidirectional atlention, and are discrete random variables in 1 ,T}
3 Linear Attention as a Graphical Model
In this section, we first show that, under certain assumptions, the vanilla linear attention mecha- nism (Katharopoulos et al., 2220)) can be expressed as an undirected graphical model. Then we introduce Latte, novel direct parameterisation, which further extend by combining it with local standard atten- tion. Our work is among the first which combines lincar attention and standard attention, such that the weights are fully normalised across the sequence.
3.1 Undirected parametrisation
The main disadvantage of attention is the quadratic time complexity which comes from the Q K matrix multiplications in Equation 1 Linear attention avoids the quadratic cost by approximating exp (xTy) y) ((z)((()) where Φ R R and making use of the associativity property of matrices.
We look at the expression only for the causal case, the bidirectional being very similar:
2. = E. T (k,)). [Ein v.c(kk)) d(q:) S.6(q:)
(4)
E.= p(q:) o(ka)
o(q)) (t o(ks) o(q.)
where S = T D' and we then recursively write S. S (kk)) with the normalisa- tion (kkk) E D and z, 2-1 + ((ke).
The basis function & can be any function, but when is positive, we can assign a probabilistic interpretation for linear attention using the Markov model with a discrete latent variable with states as in Figure 2
Figure 2: Graphical model for linear attention, discrete random variables and 1 is discrete latent variable.
3

---

Published in Transactions on Machine Learning Research (12/2025)
21
3/2
-233
-14
21
222
33
24
271
272
213
JF4
21
J/2
3/3
244
Figure 3: Token-token interaction diagram for non-causal attention of the p(s]t) matrix, see also Lin et al. (2021p. (Left) Standard Attention computes all pairwise similarities between the elements of the sequence. (Right) Latte computes only the pairwise similarities between each element of the sequence and each latent state.
In this model have:
p(s,t)
L
p(s,1,t)
((s,1))(1.t) Σ (V, E
ats = p(s(t) = p(s'))
(5)
Ei Esem 1(,t,
where we parametrised the joint ((s,(,t) ((a.))01).))
We observe that for (bt, = d(q)), (s,, = ((ks). Equation 5 is equivalent to Equation 4 This shows that for the special case when the basis function á is positive, linear attention has probabilistic interpretation. Additionally, for < T, linear attention gives low rank parametrisation of attention. Nonetheless, this parametrisation is symmetric, which does not reflect the directed nature of language.
4 Latte: Directed Parametrisation
Using the connection we have made between linear attention and the undirected latent variable model, we propose our model, Latte. Latte novel directed parameterisation for low-rank linear attention as shown in Figure 4 The directed formulation reflects asymmetric dependencies, property essential for capluring causal or temporal structure and naturally compatible with the flow of information in attention mechanisms, Our model can be seen as performing clustering around the latent states, which can be thought general learnable concepts such as shapes or colours. While attention performs the normalised similarity between each pair of tokens v, and æt, our directed latent parametrisation, performs similarity between a token æ; and the latent state
9
t
Figure 4: Graphical model for Latte, I, and s are discrete random variables.
Let / be latent variable with L possible states and .,,..., ET sequence of imput tokens. Our model takes the assumption of independence between tokens at time and t: s tll. We define the bidirectional Latte case as:
T
T Σ p( Σ p(s(1)o
p(s 10m,
(6)
s=1 1=1
1=1
A more intuitive explanation of our model can be found in Figure 3 Similarly, we define the causal case as:
t. p(s, Σ p(1)) Σ p(s|1,t)o.
(7)
We note that p(s[l,t) t) does not imply quadratic dependency like the standard attention, since we only use to define normalisation. This becomes clear in Equation g which also highlights key difference wetween
4

---

Published in Transactions on Machine Learning Research (12/2025)
our model and the linear undirected model in Equation 5
Ots p(s(t) =
p(s] t)p(1tt
(8)
1=1
(1, Ei= 18,8,1 =1 (0(1,tt)
((,,1)
(9)
Similarly to the undirected parametrisation, we define (b(, )(q:), e and (s, <(kk)) ca!ut which ensures that the potentials are positive. For the bidirectional case, we simply have to replace in Equation to T.
We highlight that the causal case can written recursively, which enables fast generation at test time. To see this we define the normalisation terms
8.
= p !
0r =
Σ uo
(10)
f=1
and write the new representation as
L
Σ p(1lt Σ
((all,1)0
(11)
1=1
um]
Σ
calut
02 Σ
(12)
Beare,
1=1
s=1
1=1
where
"m! Nt B
alut U.
(13)
Since Ue, and 0r are cumulative sums we can use the recursions
0, 01-1,1 + w. w ह Ut-1,1 20 wt
(14)
calculate 0.+1 (i.e infer the future token) directly from 0e.1. 3r, εr,, whereas standard attention requires the full sequence 21...., Fy. In this sense, Causal Latte is a recurrent model, similar essence to Recurrent Neural Networks (RNNs) and state space models (SSMs)
From Equation 14 it immediately follows that
(Gu ct al. 2021 Smith et al. 2022 Fu ct al. 2023 3, see Figure
Our recursive form admits natural "memory" interpretation: 0τ,1 summarizes past information and is updated once per time step. However, not all linear recurrent architectures fit our probabilistic framework. For example, models such MambagGu & Dao 2023) lack the normalization needed to interpret the sequence-mixing operation In contrast, our construction supports low-rank approximations within a coherent probabilistic model (while orthogonal techniques like sparse-attention implementations are outside our present scope).
probability distribution, preventing probabilistic reading of their updates.
4.1 Hybrid Model
Lincar models reduce the time complexity from quadratic to lincar, however, their performance lacks be- hind standard attention (Yang et al.) 2024). Our formulation sullers from the same problem. Although it uses latent states to represent global concepts and share long-range information across a sequence, it may not account for local information as effectively as standard attention since lacks non-lincar element-wise comparisons. Therefore combining linear attention with standard attention is natural and has been explored
in works such as Hua et al. (2022); De et. al. 2024 Different prior models in work we maintain
5

---

Published in Transactions on Machine Learning Research (12/2025)
Figure 5: Causal Latte can be written as a recursion in which the variables = [a at,t and 0g = ........ d. contain all the information roquired to form the transformed output ot.
properly normalised attention when combine local and global context by simple extension of our latent variable model.
We achieve correctly normalised attention by defining a special latent state, / = 0, allocated to standard attention. Then the new model, called Latte Macchiato, is a weighted mixture of standard attention and Causal Latte:
- =1= - 0lt) Σ po(s)()u,
(15)
+ Σ p([)) Σ p(s](,t)0.
(16)
Here (s((t) p(l] () represents standard attention; in practice to retain computational tractability use sliding window attention with window size
& t- w 1851
po(s)
(17)
otherwise
where we now define p(l[t) to ensure normalisation over all L + states, 1= O...,,
We also show that it is possible to extend our framework to obtain state-of-tho-art results. In our definition of Latte Macchiato, the quantity p(([t) depends only on the token zt, while p(s]t, () is based only on xs from the entire sequence. To encourage the latent states to capture temporal dependencies across multiple sub-words, we can use a 1D convolution size to compute these probabilities:
K
Vi
wfx
(18)
1=
e wy
eV. m
p(llt) 1 Σ e
p(s[[,t) t)
(19)
w7
Σ
eMl w
We observed that performance improves with larger convolution sizes K prompting us to also exlend to depend on all previous tokens using a linear recurrett neural network. For our experiments, we used the recurrent gated linear unit (RGLRU) layer introduced by De et al. (2024), which, compared to a convolution, is also input dependent. Note that both the convolution and recurrent layers break positional invariance, thereby eliminating the need for positional encodings in these extensions.

---

Published in Transactions on Machine Learning Research (12/2025)
5
Experiments
5.1 Synthetic Tasks
Lincar models are known for being worse than transformers in retricval capabilities Arora ct al. 2023). Hence we test the capabilities our model on the synthetic MQAR data (Arora et al. 2023) and compare with two other linear models and the standard transformer. Figure c shows that Latte performs competitively with the transformer and outperforms the other lincar models in our training set. In all the experiments, the window size of allention is 128, being smaller than the context length.
Sequence length 256, number keys:
Sequence length 512, number keys:
1.00
0.98
0.96
0.92
Standard Causal Attention
Latte
Standard Causal Attention
Mamba
0.0
Mumbe
0.90
State dimension
State dimension
Figure 6: Accuracy (ACC) on MQAR dataset for different sequence lengths and number of key-value pairs. We set the number of test examples to 10000 and train examples to 100000.
5.2 Ablations
We train small causal models OpenWebText Gooassla & Cohen 2019 for next-token prediction using shared setup across all variants for 8B tokens.
Base model. Table I reports perplexity across variants. The base model (Latte) uses Latte layers for temporal mixing, with the rest matching a standard Transformer (Vaswaní et al.] 2017). "++" variant replaces LayerNorm with RMSNorm (Zhang Sennrich 2019) and the leedforward block with a Gated Linear Unit (GLU) (Dauphin al. 2017), yielding a small performance gain.
Local attention. We extend + with K = convolution (Section 4.1 improving results with minimal parameter overhead. Adding 128-token sliding window attention (SWA) using ROPE (Su et al. 2024h yields larger performance boost at the cost of more parameters.
Recurrent variant. Replacing the convolution with a Recurrent Gated Linear Unit (RGLRU) further improves performance and slightly reduces parameters. Combining this with SWA produces Latte-R++, our best model.
Comparison to SOTA. Table compares Lable to other efficient models. To isolate the effect of Latte, we evaluate R-SWA++, which replaces Latte attention with normalized RGLRU outputs: Y = RGLRU(X), Z RMSNorm Q W,Z, K ZW, V =xW., similar to Mega (Ma et al. 2023). Latte consistently achieves the lowest perplexity despite a modest higher parameter count. We keep depth, width, and feed forward dimensions fixed, and match Griffin's SWA window size (128).
Overall, combining global latent attention with local SWA Maatt-Maccchitoo achieves performance compet- language models. Importantly, Latte can also extend pro-trained models to longer
itive with contexts via global attention (Section 5.4)

---

Published in Transactions on Machine Learning Research (12/2025)
Table Iterative improvement of Latte language modelling. SWA: Sliding Window Attention.
Model
Params.
PPL /
Latte
111M
21.88
Latte+
140M
21.56
Latte +
140M
20.26
+on-S-Aon+
153M
18.52
Latte-R++
139M
19.99
Latte-R- Saite-R-SWA++
153M
17.64
Transformer
151M
17.19
Table 2: Comparison of Lacle-Macch with other linear-scaling models on language modelling. SWA: Sliding
Window Attention
Model
Params.
PPL 4
Mega Ma et al.
2023
153M
23.75
Retnet
Sun al.
(023)
197M
21.59
H3 Fu el al. 2023
125M
21.0
RWKV
Peng al.
(2023a.
153M
18.97
Griffin De et a..
(2024]
139M
18.83
R-SWA
our
141M
18.25
Mamba Gu & Dao
(2023)
149M
17.70
GLA Yang et al. (2024)
206M
19.10
Ligth.Ati Qin el. al. (2024a)
166M
23.67
GetedDeltaNNt Yang al. 2025
190M
21.41
+
153M
17.64
Latent Collapse. common issue with latent variable models is latent collapse, where only a small subset of latent slates used. As demonstrated in Figure 7 Latte does not exhibit latent collapse, even in the absence of dropout. The figure also highlights that local attention and Latte attention are effectively used, with the probability mass distributed across various latent states. The plots generated from various heads and layers of the Laut-RGLRU-SSAA model, as detailed in Table 2 We provide plots for all layers and heads in Figure
5.3 Bidirectional Tasks
In our second of experiments, we evaluate the bidirectional version of Latte on the Long-Range Arena (LRA) benchmark (Tay ct al.] 2021). We compare it against vanilla lincar attention (as introduced by [Katharopoulos et. al.] 2020) and other linear allention variants. As shown in Table 3 bidirections Latte outperforms vanilla linear attention and performs on par with the strongest transformer- lines, though it still falls short of state-space models that are time-invariant. When we add a recurrent linear layer, performance improves substantially on the discrete tasks in LRA. However, performance on image-based (continuous) tasks remains weaker. This pattern consistent. with other models like Mamba, which also excel in discrete domains such language modelling but underperform on continuous data]] Finally, performance on discrete tasks improves even further when using bidirectional sliding window, in Latte- Matthiaaa (Latte-R-SWA. R-SWA++).
base-
1 As noted by the authors of Mamba
Gu Dao 2023
8

---

Published in Transactions on Machine Learning Research (12/2025)
35
32
21
20
20
18
17
16-
1S
12
TI
10
10
10
11 14
10 12
15
(a)
Layer 1, Head
(b) Layer 3, Head
(c) Layer 4, Head
Figure 7: Plots of p(]]t) for different layers and heads across a sequence of tokens and O,..., 16. State corresponds using standard causal windowed attention, whereas states higher than zero correspond global latent tokens. Brighter means higher probability.
Table 3: Classification accuracies for LRA dataset. We report the best test score (higher better). All Latte versions are bidirectional.
Model
ListOps
Text
Retrieval
Image
Pathfinder
Bid. Att.
36.37
64.27
57.46
42.44
71.40
Linformer
35.70
53.94
52.27
38.50
76.34
Longformer
35.63
62.85
56.89
42.22
69.71
Luna Bid.
38.01
65.74
79.55
47.47
78.89
Linear Att.
16.13
65.90
53.09
42.34
75.30
Mega-Chunk Chunk
58.76
90.19
90.97
85.80
94.41
S5
62.15
89.31
91.40
88.00
95.33
Latte
40.18
64.51
73.39
47.55
75.61
Latte-R+4
56.7
83.85
81.07
57.61
72.13
Latte-R-SWA++
61.39
85.8
87.67
70.19
73.69
5.4 Extending Pre-trained Model
Training large models from scratch is computationally expensive, even when the sequence mixing layer (like attention) has linear time and memory complexity. Recent work has demonstrated preliminary success distilling pre-trained quadratic self-attention layers into sub-quadratic layers, such Mamba (Bick et al., 2024). However, unlike Latte, these architectures (Gu & Dao 2023) significantly differ from the attention mechanisms used in standard transformers, making knowledge distillation from pre-trained transformers more complex. Other research has modified relative embeddings in standard attention to enable sequence extrapolation, but the computational cost remains quadratic (Sun et al. 2022). We use Lattc-Macchiato with SWA weights taken from a pre-trained large model and show that training only the Latte-specific weights for 1.6B tokens is sufficient. This approach enables and effective sequence length extrapolation, by bootstrapping from pro-trained open-source large language model. In our experiments, we use pre-trained 2.6B Gemma model [Gemma-Team 2024 and replace the standard attention layer with Latte-Macchiato layer of 128 long sliding window attention. The model is trained on the SlimPajama dataset (Soboleva et al. 2023). , for a single day on four 80GB A100 GPUs. In Table g we evaluate both the original Gemma model and our modified version, Gemma Macchiato, on the validation set as well as other publicly available corpora First, on sequences of length 4K, which match the
achieve desirable properties, such as global context
We the version woth copyrighted content

---

Published in Transactions on Machine Learning Research (12/2025)
training length, we find that our model's results are comparable to exceed those of the original model. When extending the sequence length to 8K and 16K tokens, our model significantly outperforms Gemma, demonstrating that excellent context extrapolation capabilities are acquired with minimal additional training steps.
Table 4: PPL & on the validation set for 4K, 8K and 16K sequences. Maccca-M initialised from Gemma and pre-trained on Slim-Pas Sim-Pajama (SP) and Tiny-Stories (TS). Unlike Gemma-MMachhine Gemma fails to generalise to longer sequences.
Gemma 8K
Gemma-Macchiato
Data
4K
16K
4K
8K
16K
SP
10.97
36.35
294.18
10.14
9.99
10.27
Pile
7.42
19.26
243.54
7.27
6.98
7.04
OWT
10.75
38.36
252.74
10.76
10.72
10.99
TS
5.45
19.15
66.61
4.26
4.34
4.30
We also check the abilities of the distilled model on standard natural language harness of multiple-choice question-answerigg Like the general trend of linear models, performance decreases especially on tasks like MMLU (Mercat et al. 2024 Zhang et al. 2024). However, our aim is not to outperform standard attention, but to provide linear global context extension method which a better alternative to sequence truncation, often when the quadratic cost. of attention becomes a limitation.
Table 5: Common Few Shot learning benchmarks. The score accuracy or normalized accuracy (t). We use sliding window of size 128.
Model
MMLU
HellaSwag
Lambada
ARC-C
ARC-E
WinoG
Piqa
BooIQA
Gemma2 2B
53.0
73.03
69.8
53.4
80.2
71.4
79.1
73.61
Gemma-Man 2B
46.8
73.11
68.29
52.9
76.9
70.6
78.7
71.48
Mamba (3B)
26.2
71.0
41.7
68.2
65.9
78.1
71.0
GLA (1.3B)
49.8
46.9
26.6
57.2
53.9
71.8
5.5
Empirical Runtime Efficiency
We benchmark the forward-pan saward-pass runtime of both the convolutional (Latte-Cov-SWAHH and recurrent (Latte-R-SWA- variants of Latte against standard attention mechanisms across varying sequence lengths and model sizes. All models use identical hyperparameterss with batch size adjusted to keep the total number of tokens constant. As shown in Figure 8b and Figure Sa both Latte variants outperform standard attention in speed. While we also include Flash Attention as reference, our JAX-based implementation is not CUDA-optimized, making direct comparison unfair. Nevertheless, at the 2.6B scale, our linear JAX implementation outperforms Flash Attention's CUDA kernel.
Given its stronger performance and modest runtime overhead, the recurrent variant (Latte-R-SWA is generally preferable to the convolutional version, likely due to its ability to capture longer-range dependencies. Although our focus has been reducing asymptotic time and memory complexity, further kernel-level optimizations could improve Latte's practical runtime efficiency.
Related Work
The literature on efficient atlention spans a wide range of techniques, broadly categorized into downsam-
pling (Jaegle et al. 2021), random patterns Zaheer et al. 2020), learnable patterns (Wang et al. 2022
Kitaev al. 2019), sparsity (Ainslic et al.] 2020 Beltagy et al.] 2020). recurrence (Dai et al., 20199] and
low-rank approximations (Wang et al. 2020 Katharopoulos et al. 2020). Some recent efforts also exploit
hardware parallelism (Qin al. 2024b: Sun et al. 2024). For comprehensive survey, (Lin et al. 2021
10

---

Published in Transactions on Machine Learning Research (12/2025)
4000
Torward Pash Time (PM
Causal Attenkion
2000
2000
1000
500
1000
10K
32K
Sequarcce Larggth
Sequence Lengiht
(a) 400M parameter model.
(b) 2.6B parameter model.
Figure 8: Runtime in milliseconds (ms) for forward passes at different sequence lengths for (a) 400M and (b) 2.6B parameter models. Standard deviations for Matte-R-SW+++ are inchuded but too small to be visible.
Efficient Attention. Shen et al. 2021) introduced a linear-time bidirectional attention mechanism similar in form to our bidirectional Latte. However, their approach lacks latent-variable interprelation and does not extend to causal attention. Moreover, their focus is on vision tasks, whereas we target language modelling and explicitly integrate both global and local attention mechanisms.
Luna. Luna (Ma et al. 2021 performs atlention between input tokens and latent tokens in bidirectional setting. While Luna shares our use of latent structures, it differs significantly in architecture and parame- terisation, especially in the causal sctting. Unlike Latte, Luna does not incorporate local attention, making it less effective for capluring both short- and long-range dependencies in language.
State-space and Hybrid Models. State-based models (Gu et al., 2021 Smith et al. 2022 replace attention with structured recurrence. Gated variants like Mamba (Gu & Dao 2023) improve performance on discrete data by introducing input-conditioned dynamics. Hybrid models such as Mega (Ma et al., 2023). Griffin (De et al. 2024), and SSM-Trans (Zuo et al. 2022) combine local attention with recurrent layers. Our method falls into this hybrid category but uniquely combines latent-variste variable-baseee based global attention with efficient local context through sliding windows. Additionally, Katte-Macchiate enables seamless extension of pre-trained language models to long contexts with linear complexity, a capability not present in most prior work.
6 Limitations
Our framework does not cover all linear-time architectures, in particular, sparse-attention methods and non- normalized recurrent sequence mixers (c.g., Mamba) fall outside its probabilistic assumptions. Nevertheless, this perspective enables new directed linear parameterisation that preserves linear time while aligning with causal sequence modelling. Finally, because our method approximates full attention, careful validation and stress testing are required before deployment in high-stakes settings such as healthcare legal applications.
7 Conclusion
We introduced latent-variable formulation of attention that scales lincarly with sequence length and can serve drop-in replacement for standard attention. While prior work has explored low-rank approxi- mations, approach is, to the best of our knowledge, the first to reinterpret linear attention latent graphical model. This perspective enables principled derivation of both bidirectional and causal variants within unified framework, leading to strong empirical results in language modeling.
Beyond this core formulation, our framework naturally integrates local sliding window attention with global latent attention. This hybrid structure not only improves performance but also offers a practical advantage: it enables the extension of pre-trained language models to much longer contexts with minimal additional training and runtime overhead.
11

---

Published in Transactions on Machine Learning Research (12/2025)
Our experiments focused on language modeling and long-range classification, but the framework is broadly applicable. Future work will explore extensions to tasks such as question answering, sequence-to-sequence generation, and multimodal modeling.
References
Joshua Ainslic, Santiago Ontanon, Chris Alberti, Vaclav Cvicck, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. ETC: Encoding Long and Structured Inputs in Transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing
(EMNLP), 2020.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Mcasuring and improving recall in efficient language models, 2023. URL
https://arzin: http::////avvvor///.//// www./oggg//////e/ee//...o.ee alt.aaa...a//.//e
Ix Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The Long- Bong-Document Transformer. arxiv preprint (rXtt:20040051500 2020.
Aviv Bick, Kevin Y. Li, Eric P. Xing, J. Zico Kolter, and Albert Gu. Transformers to ssms: Distilling quadratic knowledge to subquadratic models, 2024. URL https : //arxivv wwwpiiivvooo//aee abs///bs//200 10189
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov, Transformer-Xi XL: Attentive Language Models beyond a Fixed-Length Context, In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language Modeling with Gated Convo- lutional Networks, 2017. URL [ttp:p///aaraii.orrr//a////2..00 08083
Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murand Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Ar- naud Doucet, David Budden, Yee Whye Tch, Razvan Pascanu, Nando De Freitas, and Caglar Gulechre. Griffin: Mixing gated linear recurrences with local attention for efficient language models, 2024. URL petps://arzis org/ wrgg/ba///aeee 240..194221211
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré Hungry Hungry Hippos: Towards Language Modeling with State Space Models. Tn International Conference Learning Representations, 2023.
Gemma- Team. Gemma 2: Improving Open Language Models at a Practical Size, 2024. URL https:
7/arxiv. blrxiv..rgg/bs//408.ooo 00118
Aaron Gokaslan and Vanya Cohen. OpenWebText Corpus, 2019. URL http://Skyylion077.gitba OpenWebTextCCrpust
Albert Gu and Tri Dao. Mamba: Lincar-Time Sequence Modeling with Selective State Spaces. arXiv preprint arXiv:2312.000 2023.
Albert Gu, Karan Goel, and Christopher Re. Efficiently Modeling Long Sequences with Structured State Spaces. In International Conference on Learning Representations, 2021.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoe V. Le. Transformer Quality in Linear Time, 2022. URL https://arxiis wwwoooggaaaa.o aaab//222222 10447
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General Perception with Iterative Attention. In International Conference Machine Learning, 2021.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. In International Conference on Machine Learning. 2020.
12

---

Published in Transactions on Machine Learning Research (12/2025)
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The Efficient Transformer. In International Conference on Learning Represenlations, 2019.
Alex Krizhevsky et al. Learning Multiple Layers of Features from Tiny Images. Master's Thesis, University of Toronto, 2009.
Tianyang Lin, Yuxin Wang, Xiangvang Liu, and Xipeng Qiu. A Survey of Transformers, 2021. URL https: https://arxiv
Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear Unified Nested Attention. Advances in Neural Information Processing Systems, 2021.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian Hc, Liangke Gui, Graham Ncubig, Jonathan May, and Luke Zeitlemoyer. Mega: Moving Average Equipped Gated Atlention. In The Eleventh International Conference on Learning Representations, 2023.
Andrew Maas, Raymond E. Daly, Peter T Pham, Dan Huang, Andrew Y Ng,, and Christopher Potts. Learning Word Vectors for Sentiment Analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, 2011.
Jean Mercat, Igor Vasiljevic, Sedrick Keh, Kushal Arora, Achal Dave, Adrien Gaidon, and Thomas Kollar. Linearixing Large Language Models, 2024. URL
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jic Zhu. RWKV: Reinventing RNNs for the Transformer Era, 2023a. URL httth::////rr: org/ aat/abbb/33555 13048
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. YaRN: Efficient Context Window Exten- sion of Large Language Models, 2023b. URL
atspg/abs/233999 00071
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Lightning Attention- 2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models, 2024a. URL
http::////xxvvoo//abbbbs./ htt//abb/2401..000 2441.006888 https://arxiv
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Various Lengths, Constant Speed: Efficient Language Modeling with Lightning Attention, 2024b. URL
2405. 17381
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Nya Sutskever, et. al. Language Models are Unsupervised Multitask Learners. OpenAl blog, 2019.
Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient Attention: Atten- tion with Linear Complexitics. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, 2021.
Jimmy TH Smith, Andrew Warrington, and Scott Linderman. Simplified State Space Layers for Sequence Modeling. In International Conference on Learning Representations, 2022.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPa- jama: A 627B Token Cleaned and Deduplicated Version of RedPajama. https://cerebras. ai/blog/
of slimpajama a htt.://.w....cccaeeeee////nnnnna 627BB
2023. URL https:
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced Transformer with Rotary Position Embedding. Neurocomputing, 2024.
13

---

Published in Transactions on Machine Learning Research (12/2025)
Weigao Sun, Zhen Qin, Dong Li, Xuyang Shen, Yu Qiao, and Yiran Zhong. Linear Attention Sequence Parallelism, 2024. URL https://arxiv, http:///axxxv..o/ab///////s
Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, and Furu Wei. A Extrapoatttap Transformer, 2022. URL https://arxix: org/ abs/221.. 10554
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive Network: A Successor to Transformer for Large Language Models, 2023. URL https:
//axxiv... wwwg/bb////ggooooo.ooooo/go 2tt///337000862
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long Range Arena: A Benchmark for Efficient Transformers. In International Conference on Learning Representations, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. Advances In Neurol Information Processing Systems, 2017.
Ningning Wang, Guobing Gan, Peng Zhang, Shnai Zhang, Junqiu Wei, Qun Liu, and Xin Jiang. Cluster- Former: Neural Clustering Attention for Efficient and Effective Transformer. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics, 2022.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-Attention with Linear Complexity. arXiv preprinl ar arxi::200.. 04706, 2020.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated Linear Attention Trans- formers with Serdware-Efficien Training, 2024. URL
06635
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule, 2025. URL htttpa://arxii.orr//aa///////
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslic, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big Bird: Transformers for Longer Sequences. Advances In Neural Information Processing Systems, 2020.
Biao Zhang and Rico Sennrich. Root Mean Square Layer Normalization, 2019. URL https:
abs bass/1910.000 07467
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The Hedgehog and the Porcupine: Expressive Linear Attentions with Softmax Mimicry, 2024. URL https: wbtg/b//////00o httpp/////rv...orr/abb////00... 04347
Simiao Zuo, Xiaodong Liu, Jian Jiao, Denis Charles, Eren Manavoglu, Tuo Zhao, and Jianfong Gao. Efficient
Long Sequence Modeling via State Space Augmented Transformer, 2022. URL https: httph::////aaat htt///axvv.or//ab..a
08136
A
Experimental Details
This section describes in detail the datasets and hyperparameters used for our language modelling and classification experiments. For all our experiments we use 4 A100 GPUs.
A.1 Language Modelling
OpenWebTexx [Gokaslan & Cohen 2019) is an open-source version of the corpus used to train GPT (Radford et al. 2019) and consists of 41 GB of data extracted from 8,013,769 769 documents. We tokenize the corpus using a pre-trained Byte Pair Encoding (BPE) tokenizer with a vocabulary of 50,267 tokens. We also ensure that sequences are consistently of length 1024 by concatenating consocutive tokenized examples until this length is reached. This eliminates the need for padding, ensuring that is a confounding factor and results in more efficient computation.
14

---

Published in Transactions on Machine Learning Research (12/2025)
A.1.1 Hyperparameters
In this section, describe the ayper-parameters used in all of the language modelling experiments. Where the hyper-parameter is missing, it means that we vary it in the experiment and its value is clear from the corresponding section in the paper.
Table 6: List of hyperparameters used in the language generation task.
Hyperparameter
Value
# Layers
# Heads
Hidden Dim (D)
512
Feed Forward Dim.
2048
Latent Dim (L)
256
Local Attention Window
128
Convolution Kernel (K)
Dropout
LR
5 10-4
LR Warmup
4000 steps
LR-Decay
Linear
# #Iters.
200K
Weight Decay
0.01
Seq. Len. (2')
512
Batch Size (B)
64
Tokenizer
BPE
Embedding Type
Learned
Unroll Factor
32
Table 7: Hyperparameters for adopting Gemma our framework. All the Gemma hyperparameters are kept intact. LR is the learning rate and "#" denotes "the number of".
HyperParam.
Value
Local Attention Window
128
Latent Dim (L)
LR
0.0006
LR RR-arrmup
2000
LR-Decay
Cosine
# #Iters.
100000
Seq. Len. (T)
1096
Batch Size (B)
Tokenizer
Gemma2
A.2 LRA Dataset
This section displays the hyperparameters employed in the bidirectional experiments and provides brief description of the synthetic datasots utilized in the LRA corpus. A more comprehensive account can be found in the original paper (Tay et al. 2021
In the experiments, one layer consists of a standard transformer block where the transformation operation that gives 3, Latte or Standard Attention. For positional encoding, we use the classic transformer sinusoidal embeddings. This convention holds for both bidirectional and unidirectional problems. A complete imple- mentation be found in code repository: https: http::///ptthu...ccrrrrreeadadpp
transformer
15

---

Published in Transactions on Machine Learning Research (12/2025)
Table 8: Hyperprameters used for training on LRA. Number of latent states L specified in the result table. H=number heads, D=hidden dimension, LR=learning rate, B=batch size, W WDeweeght decay. #Lyeers denotes the number of layers which include attention approximation of attention and non-linear projections. "Embed." is the type of embedding used by the SWA.
Dataset
# LLayers
H
D
LR
B
WD
Dropout
Epochs
Embed.
ListOps
6
4
40
128
1e-3
64
0.01
0.1
50
Rope
Text
6
256
1e-3
0.05
0.1
32
Rope
Retrieval
6
4
40
128
le-4
32
0.01
0.1
20
Rope
Image
6
512
1e-3
0.05
0.1
200
Absolute
Pathfinder
6
4
256
256
1e-3
64
0.03
0.2
200
Absolute
A.2.1
ListOps
The datasct contains sequences up length 2048 of numbers from and four operators: MAX, MEAN, MEDIAN and SUM MOD. Parentheses are used to delimit the reach of each operator. The answer also digit from to which allows us to easily transform the problem into ten-way classification task.
A.2.2 Text
The Text corpus is represented by a binary classification task with long text sequences. One can easily obtain large contexts from existent datasets by tokenizing at the character level. This part of the benchmark is derived from the IMDb (Maas ct al.) 2011 movic review corpus, resulting 4K character sequenees.
A.2.3 Retrieval
This dataset tests the ability of model to predict the similarity between two long documents. Similarly to the previous corpus, it ensures long contexts through character-level tokenization, resulting in 4K tokens per document. Using a "two-tower model setup" (Tay et al.) 2021) the total sequence length becomes 8K. This is a binary classification problem, which uses accuracy as metric.
A.2.4 Image
Alongside text, images can also exhibit long-range dependencies by flattening the original image into a sequence. The Image dataset is the sequential version of Cifar10 (Krizhevsky et al.) 2009), which contains images of 10 different entities: "airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck". To obtain sequence with one input channel we apply grayscale transformation. The model needs to predict the correct entity class, given the flattened image represented as a sequence of tokens.
A.2.5 PathFinder
This part of the benchmark is also represented by images where the task is to predict whether there is a path between two points in black-and-white image. This dataset consists of 32 32 images which after flattening result in sequences of length 1024. In general larger sequences can be created by increasing the resolution. Data is tokenized similarly to the image dataset in Section A.2.4
A.2.6 PathX
This datasct version of PathFinder where the image size increased to 128 128, resulting in flattened sequences of 16384 tokens. Since all the transformer architectures fail on this dataset, we do not add it the benchmark.
16

---

Published in Transactions on Machine Learning Research (12/2025)
B Additional Experiments
B.1 Length Extrapolation
Sequence Extrapolation. Considering that long sequences improve perplexity, we also examine the model's ability to extrapolate to sequences longer than those seen during training. Specifically, we train on 5K-token sequences from the BookCorpus dataset and evaluate on sequences up to 16K tokens, As shown in Figure g our model successfully extrapolates longer sequences, whereas the performance of standard causal attention degrades the sequence length increases. Notably, Latte-RGLR--SW++ also achieves performance comparable to standard attention on sequences seen during training. We YARN (Peng et al. 2023b) relative positional encoding for standard attention and ROPE for sliding window attention in Latte. YARN is used because it helps transformers extrapolate.
RGLRU-SWAH
25.0
Stundard
Eval PPL
20.0
15.0
0
2000
4000
6000
8000
10000
12000
14000
16000
# Tokens
Figure Sequence Length Extrapolation for Attention and Latch-Macch. Both models were trained on sequences of length 5K which are extrapolated to 16K during testing.
17

---

Published in Transactions on Machine Learning Research (12/2025)
C Causal Latte Implementation
Listing 1: Scan version of Latte.
Spartial (jax. static_argaune-(.. 6))
def causal_ latte (Nq, Wk, Wv, H, x, unrall=100):
Scan implementation of latte.
B: batch size H: or heads, т: seq_len, D: hidden dsa. b: number latent states Args:
ju:jap array (DN) parameter matrices snrolt; int usrott of the toop
Ng: anp.array(2)) (DL), je:jap.array
heads
к: imp.array(ee (BTD) iuput
Returne:
jap.array (BID) transformed output sequence
uae
def accumulata args): caus, norm cunsun, prev_ carry qa_t, curr_alph, V_t, revert_nexi jap exp(-c_mx prev_ax)
add_naxi jap.oxp(curr__lph c_ax)
nora_ jap. B"LLL" norn_cuneus revert_naxi) noru_cuneun add_maxi
carry jmp.einsur einsun carry jnp jinsuinsun ("D-L,, aad_naxi,, V_t) y jap. jap.einsum ("BHL, BHLD->BAD", Qs_t nora_cunsun, carry)
caun, revert,naxi)
return ((carry, norn_cumsur, c_ax), 7)
B. Т, D X. shape
Wk. ska.shapels [-1]
V jnp-ainaum DND BTD->TBK", ("T.BBTTTT wy, X) reshapee (T, B, 8, -1)
q jup. ("DL,BID->TBL""
Ng, X) reshapee (T, B, 8, -1)
g jap.einsun""DD ("DL, ("DL BTDD>>BB"", Wk, X) resbapee (T, B, B, -1)
maxi = "jax. lax. cunmax (X, axis=0)
init_alpha jap.zeros (shape=(B, H, L
init_carry jup. zeros ((B, H. L // B, D // B)) (Q, axis=-1)
Qs
-> accumulate.
unroll=uarol,,
init= carry, init_alpha, x(0)),
xa=[Qa, K, maxi],
y = y. trtanappos 0,
return (Br D)
18


Reviews:
[{"review_id": "VnjvVbA79H", "reviewer": "Reviewer_oDNy", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper revisits linear attention in Transformers through the perspective of probabilistic graphical models. The authors first show that standard linear attention is equivalent to an undirected latent variable model. Based on this insight, they propose a new formulation of linear attention based on directed latent variable model. The new formulation enables an interpretation aligned with the causal and sequential nature of language. The authors did experiment evaluation to showcase the effectiveness of the proposed linear attention formulation.", "claims_and_evidence_yn": "Yes", "claims_explanation": "1. The authors explained the theoretical results in details and with good motivation.\n2. I think the authors supported the results with convincing evidence, however, the size of models evaluated in the paper is not convincing, parameter scales are around 100M. How would the results transfer to larger model sizes, 7B, 13B?\n3. It is also not clear how would the proposed attention formulation perform in long-context setting. Long context is a main bottleneck for quadratic attention, while linear attention has the efficiency advantages. I think new linear attention methods should be evaluated in this setting.\n4. In the Runtime Efficiency section, I wonder why the authors do not evaluate Latte directly but rather Latte-Conv-SWA++ and Latte-R-SWA++.", "audience_yn": "Yes", "audience_explanation": "I think attention architecture studies should be interesting and practical topic for TMLR's audience.", "requested_changes": "1. Can the authors illustrate more on section 2 about the probabilistic model for attention. Why would we need to view attention from this perspective and what's the motivation?\n2. Can the authors also discuss more about the relationship between Latte to other linear attention models (e.g., Mamba) in the Section 3 or 4? For example, can Mamba also be viewed as graphical model.\n3. There is also other recent improvements on linear attention that are missing in this work, e.g., Gated DeltaNet [1]\n\n[1] Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR 2025.", "broader_impact_concerns": "I do not have broader impact concern on this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ygFU9xdHqm", "reviewer": "Reviewer_Xhk1", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces Latte, a probabilistic reformulation of linear attention that interprets it as a latent variable model and adds directional structure for causal modeling. It combines global latent and local sliding-window attention (“Latte Macchiato”) and achieves performance close to standard attention while outperforming existing linear variants like Mamba and Griffin.\n\n\nStrengths\n1. The paper provides a probabilistic interpretation of linear attention as a latent graphical model.\n2. The paper shows competitive results on both synthetic and real benchmarks, including long-sequence language modeling and pre-trained model.\n3. The combination of local sliding-window attention and global latent states is empirically effective.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The experiments on language modeling and long-sequence benchmarks consistently show gains over prior linear variants. The evidence is coherent and supports what the paper claimed.", "audience_yn": "Yes", "audience_explanation": "Researchers working on efficient Transformers, state-space models, and long-context LLMs would likely find this work relevant. The probabilistic framing of linear attention and the practical hybrid design (global + local) contribute to ongoing discussions on scaling attention efficiently, even if the improvements are incremental.", "requested_changes": "None", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "niyMfmzAkT", "reviewer": "Reviewer_K1V9", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The submission proposes a unified theoretical framework for efficient, linear-time attention mechanisms based on latent probabilistic modeling. In essence, the authors introduce a latent variable model that recasts the Transformer's attention computation as an inference problem. By doing so, many existing efficient attention variants can be derived as special cases of different inference techniques or assumptions within this unified model. The paper combines this theoretical framework with new algorithmic insights and empirical evaluations. Key contributions of the paper include:\n\n- **Latent Variable Formulation of Attention:** The authors formalize softmax attention as the result of marginalizing out a set of latent variables in a probabilistic model. This provides a principled interpretation where attention weights emerge from an underlying latent structure. Exact inference in this model recovers standard softmax attention, while approximate inference yields efficient alternatives in linear time. This unifying view ties together disparate efficient attention methods under a common probabilistic lens.\n\n- **Unification of Linear-Time Attention Methods:** The paper demonstrates how several existing linear-time attention mechanisms can be seen as instances or approximations of the proposed framework. For example, methods based on random feature approximations (such as Performer-style random feature attention), low-rank projections (such as Linformer or Nyström-based attention), or locality/hashing (as in Reformer) are all shown to correspond to particular choices of latent distributions or inference schemes. This not only explains prior methods in a single language but also clarifies the relationships and differences between them.\n\n\n- **New Efficient Attention Mechanism:** Building on the unified model, the authors propose a novel linear-time attention variant (or a general algorithmic template) derived from a specific inference strategy (e.g., a variational approximation or learned latent structure). This new method is designed to retain high accuracy while scaling linearly with sequence length. The paper provides algorithmic details and pseudocode, indicating how queries and keys interact via the introduced latent variables in practice. The proposed method aims to improve over existing approaches in either accuracy, stability, or flexibility, thanks to the principled design.\n\n- **Theoretical Analysis:** The manuscript offers theoretical results to support the framework. For instance, it shows that softmax attention is recovered as a special case of their model (ensuring correctness of the formulation), and it may provide bounds or analysis on approximation error and computational complexity for the proposed linear-time method. Any assumptions needed (such as independence structures or distribution choices for the latent variables) are stated. This analysis lends credibility to the claim that the new approach is both sound and efficient, and helps delineate when the approximation will be accurate.\n\n- **Empirical Evaluation:** The authors conduct a set of experiments to validate the effectiveness of the unified approach and the new attention mechanism. They evaluate on benchmark tasks that require long-range attention, such as sequences of length beyond what standard transformers handle efficiently. Examples could include algorithmic or synthetic sequence tasks, Long Range Arena benchmarks (text, image or list operations), or language modeling on long documents. The results show that the proposed method achieves comparable or superior performance to baseline efficient attentions (and sometimes even to exact softmax attention) while using significantly less time or memory on long sequences. These experiments support the paper’s central claims about improved efficiency and competitive accuracy.\n\n**Notable strengths:** The paper’s primary strength lies in its conceptual innovation and unification. Providing a latent probabilistic model for attention is an elegant idea that offers a fresh perspective on a very active research area. This unifying view can deepen understanding: it clarifies why certain approximations work and how they relate to each other. The framework could guide the design of new attention mechanisms in a more principled way (as evidenced by the authors indeed proposing a new variant from it). The theoretical development appears solid and the manuscript demonstrates high technical rigor: key derivations are provided, and the connection to softmax attention is clearly established. Additionally, the empirical section strengthens the work, Overall, the combination of a unifying theory, a novel algorithm, and solid experiments makes the contribution significant.\n\n**Notable weaknesses:** Despite its merits, the submission has a few weaknesses. First, while the unification is intellectually appealing, the added complexity of the latent variable model might make the resulting method more complicated to implement or tune than some existing heuristics. Readers might find the framework somewhat heavy in notation and conceptual load, especially if they are primarily interested in practical implementation. Second, the empirical evaluation, although generally positive, could be limited in scope: for example, it’s not clear if the new method has been tested at large scale or on real-world tasks like full-scale language modeling or image classification with transformers. Many experiments seem to be on benchmarks of moderate scale; it remains to be seen if the benefits hold in very large models or datasets (a common standard for attention mechanisms). Additionally, there may be missing baseline comparisons; while the paper covers several known methods, one or two recent efficient attention techniques (e.g., advanced kernel-based methods or hybrid approaches) were not included in experiments or discussion. This omission makes it harder to judge the new method’s advantage in context.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Yes. The paper’s claims are generally well-supported by both theoretical arguments and empirical evidence. On the theoretical side, the authors provide clear derivations showing how softmax attention emerges from their latent model and how different linear-time approximations correspond to approximate inference methods. This grounds their unification claim in rigorous math. Key propositions or theorems are stated to ensure the reader that the framework is sound and not just intuitive hand-waving. Empirically, the authors present convincing experiments: they benchmark the proposed approach against a range of relevant baselines (including other linear attention methods and standard softmax attention when feasible), using appropriate tasks that highlight long-sequence performance.", "audience_yn": "Yes", "audience_explanation": "Yes. The topic of efficient Transformer attention is of broad interest to the machine learning community, including researchers and practitioners who deal with long sequences or resource-intensive models. Many in TMLR’s audience are looking for ways to scale up models or deploy them more efficiently, and the findings of this paper directly contribute to that goal.", "requested_changes": "**Include Additional Baseline Comparisons (Critical):** To strengthen the empirical evaluation, the authors should compare the proposed method against at least one or two missing baselines that are relevant. In particular, more recent efficient attentions (e.g., state-space models or improved kernel-based attentions) could provide useful reference points. Including these comparisons would ensure the results are comprehensive and convince readers that the new method excels across the board. If certain baselines are not included because they are conceptually different or hard to integrate, the authors should explicitly justify their exclusion.\n\n**Clarify Theoretical Derivations and Assumptions (Critical):** While the paper’s theory is generally sound, some steps in the derivation of the latent variable framework and its connection to existing methods are not entirely transparent. The authors should add more explanation or proof details for critical steps, possibly in the appendix if space is an issue. For example, the transition from the latent model to exactly recovering softmax attention should be spelled out with all necessary assumptions (e.g., independence assumptions or specific distributions). Likewise, when claiming that a certain linear attention method is a special case of the framework, it would help to explicitly show the equations or conditions under which that equivalence holds. This added clarity will make the work more accessible and convincing, especially to readers less familiar with probabilistic modeling.\n\n**Discuss Scope and Limitations of the Framework (Suggested):** The manuscript would benefit from a section or paragraph discussing what types of attention mechanisms fall under this probabilistic framework and what types might not. Currently, the focus is on unifying known linear-time approximations; however, it would be insightful for readers to understand the boundaries of this framework. For instance, does the latent model approach accommodate sparse attention patterns or local window attentions naturally, or is it primarily suited to global low-rank/kernel approximations? A candid discussion of cases where the framework might struggle, or where the latent assumptions might break down, will add balance to the paper.", "broader_impact_concerns": "The broader impact discussion focuses primarily on computational efficiency, sustainability, and democratization, which are indeed positive. Training and inference cost reductions can significantly lower the carbon footprint and make long-context models more accessible.  Maybe the authors should acknowledge that in critical applications (e.g., healthcare, legal), approximate attention may introduce subtle reliability issues. This possibility is not discussed.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)