Параллакс: параметризованное локальное линейное внимание для языкового моделирования

Аннотация

Большие языковые модели (БЯМ) стали центральной парадигмой в искусственном интеллекте, однако базовый вычислительный примитив внимания оставался структурно неизменным. Локальное Линейное Внимание (LLA) — это механизм внимания, выведенный из непараметрической статистики в рамках регрессионного подхода, применяемого на этапе тестирования. В отличие от предыдущих исследований эффективных вариантов внимания, LLA заменяет локальную константную оценку в softmax-внимании на локальную линейную оценку, что обеспечивает доказуемо превосходный компромисс смещения и дисперсии для ассоциативной памяти. Однако LLA не масштабировалось при предобучении БЯМ из-за проблем вычислительной сложности и численной устойчивости. Мы представляем Parallax — параметризованное Локальное Линейное Внимание, пригодное для масштабирования в БЯМ. Parallax исключает численный решатель из LLA и добавляет обучаемый проектор, подобный запросу, который зондирует ковариацию ключей и значений. Мы помещаем Parallax в семейство механизмов внимания, объединённых шириной окна, конструкцией зонда и аффинной структурой. Мы предлагаем алгоритм, учитывающий особенности аппаратного обеспечения, который повышает арифметическую интенсивность по сравнению с FlashAttention, переводя внимание в более вычислительно-ограниченный режим. Наш прототип ядра декодирования сравнивается или превосходит FlashAttention 2/3 при различных размерах батча и длинах контекста. Мы предобучаем Parallax на масштабах 0,6B и 1,7B и обнаруживаем последовательное улучшение перплексии на протяжении всего предобучения, причём выигрыш переносится на последующие бенчмарки. Преимущество сохраняется как при контроле по числу параметров, так и по вычислительным затратам, что демонстрирует Парето-улучшение. Мы проводим тщательные абляции по предобучению и выявляем новый феномен: Muon раскрывает потенциал Parallax. Насколько нам известно, это первая эмпирическая демонстрация сильного совместного проектирования архитектуры и оптимизатора для механизмов внимания в литературе по архитектурам.

English

Large Language Models (LLMs) have become the central paradigm in artificial intelligence, yet the core computational primitive of attention has remained structurally unchanged. Local Linear Attention (LLA) is an attention mechanism derived from nonparametric statistics in the test-time regression framework. In contrast to prior research on efficient attention variants, LLA upgrades the local constant estimate in softmax attention to a local linear estimate, yielding provably superior bias-variance tradeoffs for associative memory. However, LLA has not been scaled in LLM pretraining due to computational and numerical stability concerns. We introduce Parallax, a parameterized Local Linear Attention that is scalable for LLMs. Parallax eliminates the numerical solver in LLA and learns an extra query-like projector that probes the KV covariance. We place Parallax within a family of attention mechanisms connected by the bandwidth, the probe construction and the affine structure. We propose a hardware-aware algorithm that increases the arithmetic intensity over FlashAttention, shifting attention into a more compute bound regime. Our prototype decode kernel matches or outperforms FlashAttention 2/3 across diverse batch sizes and context lengths. We pretrain Parallax at 0.6B and 1.7B scales and find consistent perplexity improvements throughout pretraining with gains that transfer to downstream benchmarks. The advantage persists under both parameter-matched and compute-matched controls, demonstrating a Pareto improvement. We perform careful pretraining ablations and identify a novel phenomenon whereby Muon unlocks the capacity of Parallax. To our knowledge, this is the first empirical demonstration of strong architecture-optimizer codesign for attention mechanisms in the architecture research literature.