Длинноконтекстное предварительное обучение с маяковым вниманием

Аннотация

Обучение каузальных трансформеров на последовательностях экстремальной длины ограничено квадратичными затратами времени и памяти, характерными для масштабированного скалярного произведения внимания (SDPA). В данной работе мы предлагаем «Lighthouse Attention» — алгоритм иерархического внимания на основе симметричной селекции, предназначенный исключительно для обучения. Этот алгоритм оборачивает обычный SDPA и может быть легко удалён ближе к концу обучения. Наш иерархический отбор также не требует вычисления градиентов, что избавляет от необходимости реализации сложного и потенциально неэффективного обратного прохода через ядро. Наш вклад состоит из трёх частей: (i) Этап доквадратичной иерархической пред- и постобработки, выполняющий адаптивное сжатие и разжатие последовательности. (ii) Стратегия симметричного сжатия, которая одновременно объединяет запросы, ключи и значения, сохраняя при этом каузальность слева направо, что значительно улучшает параллелизм. (iii) Двухэтапный подход к обучению: мы проводим предобучение в течение большей части времени с использованием «Lighthouse Attention», а затем восстанавливаем модель полного внимания с помощью короткого этапа обучения. Мы провели предварительные эксперименты по предобучению LLM малого масштаба, которые показывают эффективность нашего метода по сравнению с обучением на полном внимании при всех прочих равных условиях: мы достигли меньшего общего времени обучения и более низкого финального значения функции потерь после восстановительной фазы. Полный код доступен по адресу: https://github.com/ighoshsubho/lighthouse-attention

English

Training causal transformers at extreme sequence lengths is bottlenecked by the quadratic time and memory of scaled dot-product attention (SDPA). In this work, we propose Lighthouse Attention, a training-only symmetrical selection-based hierarchical attention algorithm that wraps around ordinary SDPA and can be easily removed towards the end of the training. Our hierarchical selection is also gradient-free, which exempts us from dealing with a complicated and potentially inefficient backward pass kernel. Our contribution is three-fold: (i) A subquadratic hierarchical pre- and post-processing step that does adaptive compression and decompression of the sequence. (ii) A symmetrical compression strategy that pools queries, keys and values at the same time, while preserving left-to-right causality, which greatly improves parallelism. (iii) A two stage training approach which we pre-train for the majority of the time with Lighthouse Attention and recover a full attention model at the end with a short training. We run preliminary small scale LLM pre-training experiments that show the effectiveness of our method compared to full attention training with all other settings matched, where we achieve a faster total training time and lower final loss after the recovery phase. Full code is available at: https://github.com/ighoshsubho/lighthouse-attention