Várias Comprimentos, Velocidade Constante: Modelagem de Linguagem Eficiente com Lightning Attention

Apresentamos Lightning Attention, a primeira implementação de atenção linear que mantém uma velocidade de treinamento constante para vários comprimentos de sequência sob consumo de memória fixo. Devido ao problema com operações de soma cumulativa (cumsum), implementações anteriores de atenção linear não conseguem atingir sua vantagem teórica em um cenário casual. No entanto, esse problema pode ser efetivamente resolvido utilizando diferentes estratégias de cálculo de atenção para computar as diferentes partes da atenção. Especificamente, dividimos o cálculo da atenção em intra-blocos e inter-blocos e usamos o cálculo de atenção convencional para intra-blocos e truques de kernel de atenção linear para inter-blocos. Isso elimina a necessidade de cumsum no cálculo da atenção linear. Além disso, uma técnica de segmentação (tiling) é adotada tanto nos procedimentos de avanço quanto de retrocesso para aproveitar ao máximo o hardware da GPU. Para melhorar a precisão enquanto preserva a eficácia, introduzimos o TransNormerLLM (TNL), uma nova arquitetura adaptada à nossa lightning attention. Realizamos testes rigorosos em conjuntos de dados padrão e coletados por nós, com diferentes tamanhos de modelo e comprimentos de sequência. O TNL é notavelmente mais eficiente que outros modelos de linguagem. Além disso, resultados de benchmark indicam que o TNL tem desempenho equivalente aos LLMs de última geração que utilizam estruturas convencionais de transformer. O código-fonte está disponível em github.com/OpenNLPLab/TransnormerLLM.