SSA: Разреженное разреженное внимание через выравнивание выходов полного и разреженного внимания в пространстве признаков

SSA: Sparse Sparse Attention by Aligning Full and Sparse Attention Outputs in Feature Space

November 25, 2025

Авторы: Zhenyi Shen, Junru Lu, Lin Gui, Jiazheng Li, Yulan He, Di Yin, Xing Sun

cs.AI

Аннотация

Квадратичная сложность полного механизма внимания ограничивает эффективную обработку длинных контекстов в больших языковых моделях (LLM). Разреженное внимание снижает эту вычислительную стоимость, ограничивая каждый запрос подмножеством предыдущих токенов; однако подходы, не требующие дообучения, часто приводят к значительной деградации качества. Нативные методы разреженного внимания (например, NSA, MoBA) смягчают эту проблему, но демонстрируют критический парадокс: они порождают меньшую разреженность внимания, чем модели с полным вниманием, хотя и направлены на его аппроксимацию, что может ограничивать их эффективность. Мы объясняем этот парадокс дефицитом градиентных обновлений: ключевые и value-пары с низким рангом, исключенные в процессе разреженного обучения, не получают ни прямого вклада, ни обратных градиентов и, следовательно, никогда не обучаются должному подавлению. Чтобы преодолеть это ограничение, мы предлагаем SSA (Sparse Sparse Attention) — унифицированную框架 обучения, которая учитывает как разреженное, так и полное внимание и обеспечивает двунаправленное выравнивание на каждом слое. Эта конструкция сохраняет поток градиентов ко всем токенам, одновременно явно побуждая выходы разреженного внимания соответствовать их аналогам с полным вниманием, тем самым способствуя большей разреженности. В результате SSA достигает наилучших результатов при выводе как с разреженным, так и с полным вниманием на нескольких наборах данных для проверки здравого смысла. Более того, SSA позволяет моделям плавно адаптироваться к различным бюджетам разреженности: производительность последовательно улучшается по мере увеличения числа токенов, доступных для внимания, что поддерживает гибкий баланс между вычислениями и производительностью во время вывода. Наконец, мы показываем, что нативное обучение с разреженным вниманием, что удивительно, улучшает экстраполяцию на длинных контекстах за счет смягчения перераспределения значений внимания в «сток-областях», причем SSA демонстрирует наилучшую способность к экстраполяции.

English

The quadratic complexity of full attention limits efficient long-context processing in large language models (LLMs). Sparse attention mitigates this cost by restricting each query to attend to a subset of previous tokens; however, training-free approaches often lead to severe performance degradation. Native sparse-attention methods (e.g., NSA, MoBA) alleviate this issue, yet exhibit a critical paradox: they produce lower attention sparsity than full-attention models, despite aiming to approximate full attention, which may constrain their effectiveness. We attribute this paradox to gradient update deficiency: low-ranked key-value pairs excluded during sparse training receive neither forward contribution nor backward gradients, and thus never learn proper suppression. To overcome this limitation, we propose SSA (Sparse Sparse Attention), a unified training framework that considers both sparse and full attention and enforces bidirectional alignment at every layer. This design preserves gradient flow to all tokens while explicitly encouraging sparse-attention outputs to align with their full-attention counterparts, thereby promoting stronger sparsity. As a result, SSA achieves state-of-the-art performance under both sparse and full attention inference across multiple commonsense benchmarks. Furthermore, SSA enables models to adapt smoothly to varying sparsity budgets; performance improves consistently as more tokens are allowed to attend, supporting flexible compute-performance trade-offs at inference time. Finally, we show that native sparse-attention training surprisingly improves long-context extrapolation by mitigating the over-allocation of attention values in sink areas, with SSA demonstrating the strongest extrapolation capability.