Unifying Linear-Time Attention via Latent Probabilistic Modelling

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Dolga, Rares, Maystre, Lucas, Cobzarenco, Marius, Barber, David
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908684911640576
author Dolga, Rares
Maystre, Lucas
Cobzarenco, Marius
Barber, David
author_facet Dolga, Rares
Maystre, Lucas
Cobzarenco, Marius
Barber, David
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2402_17512
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Unifying Linear-Time Attention via Latent Probabilistic Modelling
Dolga, Rares
Maystre, Lucas
Cobzarenco, Marius
Barber, David
Computation and Language
Machine Learning
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.
title Unifying Linear-Time Attention via Latent Probabilistic Modelling
topic Computation and Language
Machine Learning
url https://arxiv.org/abs/2402.17512