Unifying Linear-Time Attention via Latent Probabilistic Modelling

2024-02-27

AI Summary

To address the scalability bottleneck of quadratic-complexity attention in Transformers for long-sequence modeling—and the performance degradation and causal-structure modeling difficulties of existing linear attention methods in discrete language modeling—this paper proposes a linear-time attention framework grounded in directed probabilistic graphical models. Our method introduces (1) directed parametric linear attention, which jointly couples global latent variables with local standard attention, and (2) a recurrent query/key generation mechanism that explicitly encodes temporal causality without requiring relative positional encoding. Evaluated on standard language modeling benchmarks, our approach matches native Transformer performance while substantially outperforming diverse linear attention variants. Crucially, it retains strict O(n) time complexity, achieving a favorable trade-off between modeling fidelity and structural awareness for discrete sequences.

Technology Category

• Natural Language Processing: (Large) Language Models
• Planning, Routing, and Scheduling: Planning with Language Models
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Develops a directed linear attention model for language
• Integrates global latent-variable and local standard attention
• Achieves competitive performance without relative positional encodings

Innovation

• Directed linear attention via probabilistic graphical models
• Asymmetric structure for causal sequential language modeling
• Recurrent queries and keys without positional encodings