Atención Lineal de Orden Superior

Resumen

El costo cuadrático de la atención de producto escalar escalado es un obstáculo central para escalar modelos de lenguaje autorregresivos a contextos largos. La atención de tiempo lineal y los Modelos de Espacio de Estado (SSM) ofrecen alternativas escalables, pero normalmente están restringidos a aproximaciones de primer orden o basadas en kernels, lo que puede limitar la expresividad. Introducimos Higher-order Linear Attention (HLA), un mecanismo causal y en flujo continuo que materializa interacciones de orden superior mediante estadísticos suficientes de prefijo compactos. En el caso de segundo orden, HLA mantiene un estado de tamaño constante y calcula las salidas por token en tiempo lineal sin materializar ninguna matriz de tamaño n por n. Proporcionamos identidades de flujo continuo en forma cerrada, una variante estrictamente causal enmascarada que utiliza dos resúmenes adicionales, y un esquema de entrenamiento paralelo por bloques basado en escaneos asociativos que reproduce exactamente las activaciones de una recurrencia serial. Además, esbozamos extensiones a tercer orden y órdenes superiores. Colectivamente, estos resultados posicionan a HLA como un bloque de construcción escalable y fundamentado que combina la mezcla dependiente de datos, similar a la atención, con la eficiencia de las arquitecturas recurrentes modernas. Página del proyecto: 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.