Il Riccio & il Porcospino: Attenzioni Lineari Espressive con Mimetismo Softmax

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

February 6, 2024

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

cs.AI

Abstract

Le attenzioni lineari hanno dimostrato potenziale nel migliorare l'efficienza dei Transformer, riducendo la complessità quadratica dell'attenzione a lineare rispetto alla lunghezza della sequenza. Ciò apre prospettive entusiasmanti per (1) l'addestramento di Transformer lineari da zero, (2) la "conversione fine-tuned" di Transformer specifici per un compito in versioni lineari che recuperano le prestazioni del compito, e (3) la "conversione pretrained" di Transformer come i grandi modelli linguistici in versioni lineari fine-tunabili su compiti downstream. Tuttavia, le attenzioni lineari spesso ottengono prestazioni inferiori rispetto all'attenzione softmax standard in termini di qualità. Per colmare questo divario di prestazioni, abbiamo osservato che le precedenti attenzioni lineari mancano di proprietà chiave dell'attenzione softmax legate a buone prestazioni: pesi a bassa entropia (o "spiky") e monotonicità del prodotto scalare. Abbiamo inoltre osservato mappe di caratteristiche sorprendentemente semplici che mantengono queste proprietà e raggiungono le prestazioni della softmax, ma sono inefficienti da calcolare nell'attenzione lineare. Proponiamo quindi Hedgehog, un'attenzione lineare apprendibile che mantiene le proprietà spiky e monotoniche dell'attenzione softmax pur conservando una complessità lineare. Hedgehog utilizza semplici MLP addestrabili per produrre pesi di attenzione che imitano l'attenzione softmax. Gli esperimenti mostrano che Hedgehog recupera oltre il 99% della qualità del Transformer standard in contesti di addestramento da zero e conversione fine-tuned, superando le precedenti attenzioni lineari fino a 6 punti di perplexity su WikiText-103 con GPT causali, e fino a 8,7 punti di punteggio GLUE su BERT bidirezionali fine-tuned. Hedgehog abilita anche la conversione pretrained. Convertire un GPT-2 pretrained in una variante di attenzione lineare raggiunge uno stato dell'arte di 16,7 di perplexity su WikiText-103 per modelli decoder subquadratici da 125M. Infine, trasformiamo un Llama-2 7B pretrained in un Llama con attenzione lineare praticabile. Con l'adattamento a basso rango, Hedgehog-Llama2 7B raggiunge 28,1 punti ROUGE-1 in più rispetto al modello base con attenzione standard, mentre le precedenti attenzioni lineari portano a un calo di 16,5 punti.

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.