Полное внимание наносит ответный удар: перенос полного внимания в разреженное за сто шагов обучения

Аннотация

Вывод с длинным контекстом в больших языковых моделях ограничен квадратичной стоимостью полного внимания. Существующие эффективные альтернативы часто опираются либо на изначально разреженное обучение, либо на эвристическое удаление токенов, создавая нежелательный компромисс между эффективностью, стоимостью обучения и точностью. В данной работе мы показываем, что LLM с полным вниманием уже внутренне разрежены и могут быть преобразованы в сильно разреженные модели с минимальной адаптацией. Наш подход основан на трех наблюдениях: (1) лишь небольшая подгруппа голов внимания действительно требует полной обработки длинного контекста; (2) долгосрочное извлечение управляется в первую очередь низкоразмерным подпространством, что позволяет эффективно извлекать релевантные токены с помощью 16-мерного индексатора; (3) полезный бюджет токенов сильно зависит от запроса, что делает динамический выбор top-p более подходящим, чем фиксированное top-k разрежение. Основываясь на этих идеях, мы предлагаем RTPurbo, который сохраняет полный кэш KV только для голов извлечения и вводит легковесный индексатор токенов для разреженного внимания. Используя внутреннюю разреженность модели, RTPurbo достигает разрежения всего за несколько сотен шагов обучения. Эксперименты на эталонных тестах с длинным контекстом и задачах рассуждения показывают, что RTPurbo сохраняет точность, близкую к безошибочной, обеспечивая при этом значительный прирост эффективности, включая ускорение предзаполнения до 9,36 раз при контексте в 1М и ускорение декодирования примерно в 2,01 раза. Эти результаты свидетельствуют о том, что сильное разреженное умозаключение может быть получено из стандартного обучения с полным вниманием без дорогостоящего собственного разреженного предобучения.

English

Long-context inference in large language models is bottlenecked by the quadratic cost of full attention. Existing efficient alternatives often rely either on native sparse training or on heuristic token eviction, creating an undesirable trade-off among efficiency, training cost, and accuracy. In this work, we show that full-attention LLMs are already intrinsically sparse and can be transformed into highly sparse models with only minimal adaptation. Our approach is built on three observations: (1) only a small subset of attention heads truly requires full long-context processing; (2) long-range retrieval is governed primarily by a low-dimensional subspace, allowing relevant tokens to be retrieved efficiently with a 16-dimensional indexer; and (3) the useful token budget is strongly query-dependent, making dynamic top-p selection more suitable than fixed top-k sparsification. Based on these insights, we propose RTPurbo, which retains the full KV cache only for retrieval heads and introduces a lightweight token indexer for sparse attention. By exploiting the model's intrinsic sparsity, RTPurbo achieves sparsification with only a few hundred training steps. Experiments on long-context benchmarks and reasoning tasks show that RTPurbo preserves near-lossless accuracy while delivering substantial efficiency gains, including up to a 9.36times prefill speedup at 1M context and about a 2.01times decode speedup. These results suggest that strong sparse inference can be obtained from standard full-attention training without expensive native sparse pretraining.