Высокоуровневое линейное внимание

Аннотация

Квадратичная стоимость масштабированного скалярно-продуктового внимания является ключевым препятствием для масштабирования авторегрессионных языковых моделей на длинные контексты. Линейные по времени механизмы внимания и модели пространства состояний (SSM) предлагают масштабируемые альтернативы, но обычно ограничиваются аппроксимациями первого порядка или на основе ядер, что может ограничивать выразительность. Мы представляем Высокоуровневое Линейное Внимание (HLA) — каузальный потоковый механизм, реализующий взаимодействия высшего порядка через компактные достаточные статистики префикса. В случае второго порядка HLA поддерживает состояние постоянного размера и вычисляет выходы для каждого токена за линейное время без материализации каких-либо матриц размера n × n. Мы приводим замкнутые потоковые тождества, строго каузальную маскированную версию с использованием двух дополнительных сводок и схему параллельного обучения на чанках, основанную на ассоциативных сканированиях, которая точно воспроизводит активации последовательной рекуррентности. Мы также намечаем расширения для третьего и более высоких порядков. В совокупности эти результаты позиционируют HLA как принципиально масштабируемый строительный блок, сочетающий зависящее от данных смешивание, подобное вниманию, с эффективностью современных рекуррентных архитектур. Страница проекта: https://github.com/yifanzhang-pro/HLA.

English

The quadratic cost of scaled dot-product attention is a central obstacle to scaling autoregressive language models to long contexts. Linear-time attention and State Space Models (SSMs) provide scalable alternatives but are typically restricted to first-order or kernel-based approximations, which can limit expressivity. We introduce Higher-order Linear Attention (HLA), a causal, streaming mechanism that realizes higher interactions via compact prefix sufficient statistics. In the second-order case, HLA maintains a constant-size state and computes per-token outputs in linear time without materializing any n times n matrices. We give closed-form streaming identities, a strictly causal masked variant using two additional summaries, and a chunk-parallel training scheme based on associative scans that reproduces the activations of a serial recurrence exactly. We further outline extensions to third and higher orders. Collectively, these results position HLA as a principled, scalable building block that combines attention-like, data-dependent mixing with the efficiency of modern recurrent architectures. Project Page: https://github.com/yifanzhang-pro/HLA.