TriAttention: Эффективное длинное рассуждение с тригонометрической компрессией ключей и значений

TriAttention: Efficient Long Reasoning with Trigonometric KV Compression

April 6, 2026

Авторы: Weian Mao, Xi Lin, Wei Huang, Yuxin Xie, Tianfu Fu, Bohan Zhuang, Song Han, Yukang Chen

cs.AI

Аннотация

Расширенное рассуждение в больших языковых моделях (LLM) создает серьезные проблемы с пропускной способностью памяти KV-кэша. Ведущие методы сжатия KV-кэша оценивают важность ключей и значений (KV) с помощью оценок внимания из недавних запросов после применения RoPE. Однако запросы поворачиваются вместе с позицией во время RoPE, что делает репрезентативные запросы крайне малочисленными, что приводит к плохому выбору топ-ключей и нестабильному рассуждению. Чтобы избежать этой проблемы, мы обращаемся к пространству до применения RoPE, где наблюдаем, что векторы Q и K сильно сконцентрированы вокруг фиксированных ненулевых центров и остаются стабильными across позиций — Концентрация Q/K. Мы показываем, что эта концентрация приводит к тому, что запросы предпочтительно обращают внимание на ключи на определенных расстояниях (например, ближайшие ключи), причем центры определяют, какие расстояния предпочтительны, через тригонометрический ряд. Основываясь на этом, мы предлагаем TriAttention для оценки важности ключей, используя эти центры. Через тригонометрический ряд мы используем предпочтение расстояния, характеризуемое этими центрами, для оценки ключей в соответствии с их позициями, а также используем нормы Q/K в качестве дополнительного сигнала для оценки важности. На AIME25 с генерацией 32K токенов TriAttention соответствует точности рассуждения Full Attention, достигая при этом в 2.5 раза более высокой пропускной способности или сокращения памяти KV в 10.7 раз, тогда как ведущие базовые методы достигают лишь около половины точности при той же эффективности. TriAttention позволяет развернуть OpenClaw на одном потребительском GPU, где длинный контекст в противном случае вызывал бы переполнение памяти с Full Attention.

English

Extended reasoning in large language models (LLMs) creates severe KV cache memory bottlenecks. Leading KV cache compression methods estimate KV importance using attention scores from recent post-RoPE queries. However, queries rotate with position during RoPE, making representative queries very few, leading to poor top-key selection and unstable reasoning. To avoid this issue, we turn to the pre-RoPE space, where we observe that Q and K vectors are highly concentrated around fixed non-zero centers and remain stable across positions -- Q/K concentration. We show that this concentration causes queries to preferentially attend to keys at specific distances (e.g., nearest keys), with the centers determining which distances are preferred via a trigonometric series. Based on this, we propose TriAttention to estimate key importance by leveraging these centers. Via the trigonometric series, we use the distance preference characterized by these centers to score keys according to their positions, and also leverage Q/K norms as an additional signal for importance estimation. On AIME25 with 32K-token generation, TriAttention matches Full Attention reasoning accuracy while achieving 2.5x higher throughput or 10.7x KV memory reduction, whereas leading baselines achieve only about half the accuracy at the same efficiency. TriAttention enables OpenClaw deployment on a single consumer GPU, where long context would otherwise cause out-of-memory with Full Attention.