Unifying Linear-Time Attention via Latent Probabilistic Modelling

TL;DR

Quadratic attention limits hinder long-sequence modelling; this paper reinterprets linear attention through probabilistic graphical models, revealing an undirected latent-variable view and introducing a directed Latte parameterisation. It then presents Latte-Macchiato, a hybrid Latent global plus local attention with full normalization, plus a recurrent formulation of queries and keys to avoid relative positional encodings. The approach yields competitive language modelling performance with existing linear variants and enables efficient extension of pre-trained models to longer contexts. Empirical results on synthetic tasks and LRA show favorable trade-offs between accuracy and speed, supporting practicality for long-context NLP.

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.

Figures (9)

• Figure 1: Graphical model for bidirectional attention, $t$ and $s$ are discrete random variables in $\{0, \cdots ,T\}$
• Figure 2: Graphical model for linear attention, $t$, $s$ are discrete random variables and $l$ is a discrete latent variable.
• Figure 3: Token-token interaction diagram for non-causal attention of the $p(s|t)$ matrix, see also lin2021surveytransformers. (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.
• Figure 4: Graphical model for Latte, $l$, $t$ and $s$ are discrete random variables.
• Figure 5: Causal Latte can be written as a recursion in which the variables $\alpha_t=\left[{\alpha_{t,1},\ldots, \alpha_{t,L}}\right]$ and $\tilde{v}_t=\left[{\tilde{v}_{t,1},\ldots, \tilde{v}_{t,L}}\right]$ contain all the information required to form the transformed output $\tilde{x}_t$.