El Erizo y el Puercoespín: Atenciones Lineales Expresivas con Imitación 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

Resumen

Las atenciones lineales han demostrado potencial para mejorar la eficiencia de los Transformers, reduciendo la complejidad cuadrática de la atención a lineal en la longitud de la secuencia. Esto ofrece perspectivas emocionantes para (1) entrenar Transformers lineales desde cero, (2) la "conversión ajustada" de Transformers específicos para tareas en versiones lineales que recuperen el rendimiento de la tarea, y (3) la "conversión preentrenada" de Transformers, como modelos de lenguaje grandes, en versiones lineales ajustables para tareas posteriores. Sin embargo, las atenciones lineales a menudo tienen un rendimiento inferior al de la atención softmax estándar en términos de calidad. Para cerrar esta brecha de rendimiento, encontramos que las atenciones lineales previas carecen de propiedades clave de la atención softmax vinculadas a un buen rendimiento: pesos de baja entropía (o "puntiagudos") y monotonicidad del producto escalar. Además, observamos mapas de características sorprendentemente simples que conservan estas propiedades y coinciden con el rendimiento de softmax, pero son ineficientes de calcular en la atención lineal. Por lo tanto, proponemos Hedgehog, una atención lineal aprendible que conserva las propiedades puntiagudas y monótonas de la atención softmax mientras mantiene la complejidad lineal. Hedgehog utiliza MLPs entrenables simples para producir pesos de atención que imitan la atención softmax. Los experimentos muestran que Hedgehog recupera más del 99% de la calidad del Transformer estándar en configuraciones de entrenamiento desde cero y conversión ajustada, superando a las atenciones lineales previas hasta en 6 puntos de perplejidad en WikiText-103 con GPTs causales, y hasta en 8.7 puntos de puntuación GLUE en BERTs bidireccionales ajustados. Hedgehog también permite la conversión preentrenada. Convertir un GPT-2 preentrenado en una variante de atención lineal logra un estado del arte de 16.7 de perplejidad en WikiText-103 para modelos decodificadores subcuadráticos de 125M. Finalmente, convertimos un Llama-2 7B preentrenado en un Llama viable con atención lineal. Con adaptación de bajo rango, Hedgehog-Llama2 7B logra 28.1 puntos ROUGE-1 más altos que el modelo base de atención estándar, donde las atenciones lineales previas provocan caídas de 16.5 puntos.

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.