O Ouriço e o Porco-espinho: Atenções Lineares Expressivas com Mimetismo de Softmax

The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry

February 6, 2024

Autores: Michael Zhang, Kush Bhatia, Hermann Kumbong, Christopher Ré

cs.AI

Resumo

Atenções lineares têm demonstrado potencial para melhorar a eficiência dos Transformers, reduzindo a complexidade quadrática da atenção para linear em relação ao comprimento da sequência. Isso traz uma promessa empolgante para (1) treinar Transformers lineares do zero, (2) "conversão ajustada" de Transformers específicos para tarefas em versões lineares que recuperam o desempenho da tarefa, e (3) "conversão pré-treinada" de Transformers, como grandes modelos de linguagem, em versões lineares ajustáveis para tarefas subsequentes. No entanto, as atenções lineares frequentemente têm desempenho inferior ao da atenção softmax padrão em termos de qualidade. Para fechar essa lacuna de desempenho, descobrimos que as atenções lineares anteriores carecem de propriedades-chave da atenção softmax associadas a um bom desempenho: pesos de baixa entropia (ou "pontiagudos") e monotonicidade do produto escalar. Além disso, observamos mapas de características surpreendentemente simples que retêm essas propriedades e igualam o desempenho da softmax, mas são ineficientes para calcular na atenção linear. Assim, propomos o Hedgehog, uma atenção linear aprendível que mantém as propriedades pontiagudas e monotônicas da atenção softmax, preservando a complexidade linear. O Hedgehog usa MLPs simples e treináveis para produzir pesos de atenção que imitam a atenção softmax. Experimentos mostram que o Hedgehog recupera mais de 99% da qualidade do Transformer padrão em configurações de treinamento do zero e conversão ajustada, superando as atenções lineares anteriores em até 6 pontos de perplexidade no WikiText-103 com GPTs causais, e até 8,7 pontos no GLUE score em BERTs bidirecionais ajustados. O Hedgehog também permite a conversão pré-treinada. Converter um GPT-2 pré-treinado em uma variante de atenção linear alcança o estado da arte com 16,7 de perplexidade no WikiText-103 para modelos decodificadores subquadráticos de 125M. Por fim, transformamos um Llama-2 7B pré-treinado em um Llama viável com atenção linear. Com adaptação de baixo posto, o Hedgehog-Llama2 7B alcança 28,1 pontos ROUGE-1 a mais em relação ao modelo base de atenção padrão, enquanto as atenções lineares anteriores resultam em quedas de 16,5 pontos.

English

Linear attentions have shown potential for improving Transformer efficiency, reducing attention's quadratic complexity to linear in sequence length. This holds exciting promise for (1) training linear Transformers from scratch, (2) "finetuned-conversion" of task-specific Transformers into linear versions that recover task performance, and (3) "pretrained-conversion" of Transformers such as large language models into linear versions finetunable on downstream tasks. However, linear attentions often underperform standard softmax attention in quality. To close this performance gap, we find prior linear attentions lack key properties of softmax attention tied to good performance: low-entropy (or "spiky") weights and dot-product monotonicity. We further observe surprisingly simple feature maps that retain these properties and match softmax performance, but are inefficient to compute in linear attention. We thus propose Hedgehog, a learnable linear attention that retains the spiky and monotonic properties of softmax attention while maintaining linear complexity. Hedgehog uses simple trainable MLPs to produce attention weights mimicking softmax attention. Experiments show Hedgehog recovers over 99% of standard Transformer quality in train-from-scratch and finetuned-conversion settings, outperforming prior linear attentions up to 6 perplexity points on WikiText-103 with causal GPTs, and up to 8.7 GLUE score points on finetuned bidirectional BERTs. Hedgehog also enables pretrained-conversion. Converting a pretrained GPT-2 into a linear attention variant achieves state-of-the-art 16.7 perplexity on WikiText-103 for 125M subquadratic decoder models. We finally turn a pretrained Llama-2 7B into a viable linear attention Llama. With low-rank adaptation, Hedgehog-Llama2 7B achieves 28.1 higher ROUGE-1 points over the base standard attention model, where prior linear attentions lead to 16.5 point drops.