Nous Research apresenta Lighthouse Attention para acelerar o treinamento de LLMs

A Nous Research publicou o Lighthouse Attention, um novo método para otimizar o treinamento de grandes modelos de linguagem em contextos longos. O mecanismo atua exclusivamente durante o pré-treinamento e é totalmente desativado quando esse processo termina, sem afetar a arquitetura nem o comportamento do modelo pronto no forward pass.

Como o Lighthouse Attention funciona

O Lighthouse Attention é um mecanismo de atenção hierárquica e seletiva que envolve a atenção padrão de produto escalar escalonado (scaled dot-product attention) durante o pré-treinamento do modelo. Na prática, a cada passagem por uma camada de atenção o modelo usa um mecanismo seletivo específico, em vez da atenção completa sobre todos os tokens do contexto.

A principal diferença em relação a abordagens anteriores (como NSA e HISA) é o pooling simétrico de todos os componentes da atenção. Os métodos anteriores reduziam apenas chaves e valores (K e V) e ignoravam as consultas, enquanto o Lighthouse aplica pooling a consultas, chaves e valores (Q, K e V) ao mesmo tempo, por meio de uma pirâmide de resoluções em vários níveis. Isso resulta em uma redução de computação mais equilibrada e eficiente em todos os níveis.

Tecnicamente, isso reduz a complexidade computacional da operação de atenção de O(N·S·d) para O(S²·d), em que N é o comprimento total do contexto, S é o tamanho da subsequência compacta selecionada e d é a dimensão da camada oculta do modelo. Após a seleção, o FlashAttention padrão roda sobre a pequena subsequência densa, o que economiza bastante tanto recursos de computação quanto a memória de GPU necessária.

Resultados impressionantes

A Nous Research testou o Lighthouse Attention em um modelo estilo Llama-3 de 530 milhões de parâmetros com contexto de 98 mil tokens, que já é um contexto bastante longo para testes. Os resultados mostraram melhora significativa e consistente no desempenho do treinamento:

• Aceleração de 1.40–1.69 vez no treinamento de ponta a ponta em relação à implementação de referência cuDNN SDPA em GPU
• Perda de treinamento (training loss) final comparável ou menor, o que indica ausência de perda de qualidade e de precisão do modelo
• Compatibilidade total com a infraestrutura existente do FlashAttention e com frameworks padrão como o PyTorch

Isso significa que as organizações poderiam treinar modelos grandes de 40 a 70 por cento mais rápido, sem comprometer a qualidade nem a precisão. Para modelos grandes treinados em conjuntos de dados enormes, isso se traduz em semanas de tempo de computação economizadas em clusters de GPU caros.

Aplicação prática e escalabilidade

As principais vantagens do Lighthouse Attention são a facilidade de adoção e a ausência de influência sobre o comportamento do modelo pronto. O mecanismo é usado apenas durante o pré-treinamento e se desativa automaticamente ao fim dessa etapa crítica. Por isso, um modelo treinado com Lighthouse é totalmente compatível com aplicativos, serviços e fluxos de trabalho existentes, sem mudanças em código, infraestrutura ou implantação.

A aceleração é especialmente valiosa para organizações que treinam modelos grandes com contextos de dezenas e centenas de milhares de tokens. Entre os usos típicos estão a análise de documentos e relatórios longos, a busca em texto completo em grandes bases de conhecimento, a escrita e a análise de código com contextos de 100K+ tokens, o processamento de conversas com longo histórico de mensagens e o trabalho com artigos científicos e patentes. Cada ponto percentual de computação economizado significa economia concreta de energia elétrica e uma redução considerável na conta de computação em nuvem.

O que isso significa para a pesquisa e a indústria

A otimização do treinamento de transformers continua sendo uma área de pesquisa ativa e fértil, apesar de vinte anos de investimento nos mecanismos centrais da arquitetura. O Lighthouse Attention mostra que, mesmo em arquiteturas de atenção bem estudadas e lapidadas, ainda há espaço para inovação, melhorias e otimizações inesperadas.

Se a comunidade de pesquisa adotar métodos como esse e eles forem incorporados em larga escala a frameworks abertos populares, como PyTorch, HuggingFace Transformers e outros, isso pode reduzir de forma significativa a barreira de entrada para organizações, startups e grupos de pesquisa que queiram treinar seus próprios grandes modelos de linguagem sem precisar de enormes recursos computacionais e orçamentos.