Nous Research выпустила Lighthouse Attention, ускорение предобучения длинных последовательностей в 1.4-1.7 раза

AIMPACT сообщение, 17 мая (UTC+8), Nous Research представила метод Lighthouse Attention, который с помощью селективного иерархического механизма внимания решает проблему квадратичного роста затрат на вычисление внимания при обучении на длинных последовательностях. Этот метод выполняет симметричное объединение Query, Key и Value, при этом логика выбора размещена вне ядра внимания, что позволяет использовать переиспользуемое ядро FlashAttention, и применяет двухэтапную стратегию обучения. В реальных тестах на NVIDIA B200 при длине контекста около 512K скорость прямого распространения увеличилась в 21 раз, а совокупная скорость прямого и обратного распространения — в 17.3 раза, при этом пропускная способность первого этапа достигла 126 тысяч токенов/с/графический процессор (по сравнению с 46 тысячами у плотного SDPA), а ускорение от конца до конца составило от 1.40× до 1.69×, при этом сохраняется сопоставимый или более низкий уровень потерь при обучении. Проверка на модели в стиле Llama-3 с 530M параметров показала, что три варианта Lighthouse дают конечные потери (0.698–0.71), которые лучше, чем у обученной с нуля плотной SDPA-базовой модели (0.7237), что экономит от 22.5 до 27 часов времени обучения. Статья arXiv:2605.06554.