Der Igel & das Stachelschwein: Ausdrucksstarke lineare Aufmerksamkeitsmechanismen mit Softmax-Nachahmung

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

February 6, 2024

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

cs.AI

Zusammenfassung

Lineare Attention-Mechanismen haben Potenzial gezeigt, die Effizienz von Transformern zu verbessern, indem sie die quadratische Komplexität der Attention auf eine lineare Abhängigkeit von der Sequenzlänge reduzieren. Dies eröffnet vielversprechende Möglichkeiten für (1) das Training linearer Transformer von Grund auf, (2) die „feinabgestimmte Konvertierung“ von aufgaben-spezifischen Transformern in lineare Versionen, die die Aufgabenleistung wiederherstellen, und (3) die „vorab trainierte Konvertierung“ von Transformern wie großen Sprachmodellen in lineare Versionen, die für nachgelagerte Aufgaben feinabgestimmt werden können. Allerdings schneiden lineare Attention-Mechanismen oft qualitativ schlechter ab als die Standard-Softmax-Attention. Um diese Leistungslücke zu schließen, stellen wir fest, dass frühere lineare Attention-Mechanismen Schlüsseleigenschaften der Softmax-Attention, die mit guter Leistung verbunden sind, vermissen: niedrige Entropie (oder „spitze“) Gewichte und Monotonie des Skalarprodukts. Wir beobachten weiterhin überraschend einfache Feature-Maps, die diese Eigenschaften beibehalten und die Leistung der Softmax-Attention erreichen, aber ineffizient in der linearen Attention zu berechnen sind. Daher schlagen wir Hedgehog vor, einen lernbaren linearen Attention-Mechanismus, der die spitzen und monotonen Eigenschaften der Softmax-Attention beibehält, während er die lineare Komplexität bewahrt. Hedgehog verwendet einfache trainierbare MLPs, um Attention-Gewichte zu erzeugen, die die Softmax-Attention nachahmen. Experimente zeigen, dass Hedgehog über 99 % der Qualität des Standard-Transformers in Trainings- und Feinabstimmungsszenarien wiederherstellt und frühere lineare Attention-Mechanismen um bis zu 6 Perplexity-Punkte auf WikiText-103 mit kausalen GPTs und um bis zu 8,7 GLUE-Score-Punkte bei feinabgestimmten bidirektionalen BERTs übertrifft. Hedgehog ermöglicht auch die vorab trainierte Konvertierung. Die Konvertierung eines vorab trainierten GPT-2 in eine lineare Attention-Variante erreicht einen state-of-the-art Perplexity-Wert von 16,7 auf WikiText-103 für 125M subquadratische Decoder-Modelle. Schließlich wandeln wir einen vorab trainierten Llama-2 7B in einen funktionsfähigen linearen Attention-Llama um. Mit Low-Rank-Adaptation erreicht Hedgehog-Llama2 7B 28,1 höhere ROUGE-1-Punkte gegenüber dem Basis-Standard-Attention-Modell, während frühere lineare Attention-Mechanismen zu einem Rückgang von 16,5 Punkten führen.

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.