Pré-treinamento de Contexto Longo com Atenção Farol

Resumo

O treinamento de transformadores causais em comprimentos de sequência extremos é limitado pelo tempo e memória quadráticos da atenção de produto escalar escalado (SDPA). Neste trabalho, propomos a Atenção Lighthouse, um algoritmo de atenção hierárquica baseado em seleção simétrica exclusivo para treinamento, que envolve a SDPA comum e pode ser facilmente removido próximo ao final do treinamento. Nossa seleção hierárquica também é livre de gradientes, o que nos isenta de lidar com um kernel de retropropagação complicado e potencialmente ineficiente. Nossa contribuição é tripla: (i) Uma etapa de pré e pós-processamento hierárquico subquadrático que realiza compressão e descompressão adaptativas da sequência. (ii) Uma estratégia de compressão simétrica que agrupa consultas, chaves e valores simultaneamente, preservando a causalidade da esquerda para a direita, o que melhora significativamente o paralelismo. (iii) Uma abordagem de treinamento em duas etapas, na qual pré-treinamos a maior parte do tempo com a Atenção Lighthouse e recuperamos um modelo de atenção completa no final com um treinamento curto. Realizamos experimentos preliminares de pré-treinamento de LLM em pequena escala que mostram a eficácia do nosso método em comparação com o treinamento de atenção completa com todas as outras configurações equivalentes, onde obtemos um tempo total de treinamento mais rápido e uma perda final menor após a fase de recuperação. O código completo está disponível em: 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