Atención de Ranura con Compuertas para Modelado Eficiente de Secuencias en Tiempo Lineal

Gated Slot Attention for Efficient Linear-Time Sequence Modeling

September 11, 2024

Autores: Yu Zhang, Songlin Yang, Ruijie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, Peng Zhou, Guohong Fu

cs.AI

Resumen

Los Transformers de atención lineal y sus variantes con compuertas, conocidos por permitir el entrenamiento paralelo y una inferencia recurrente eficiente, aún se quedan cortos en tareas intensivas en recuperación en comparación con los Transformers tradicionales y requieren recursos significativos para el entrenamiento desde cero. Este artículo presenta la Atención con Ranuras con Compuertas (GSA), que mejora la Atención con Control de Memoria Limitada (ABC) al incorporar un mecanismo de compuerta inspirado en la Atención Lineal con Compuertas (GLA). Esencialmente, GSA consta de dos capas de GLA conectadas a través de una función softmax, utilizando lectura de memoria consciente del contexto y olvido adaptativo para mejorar la capacidad de memoria manteniendo un tamaño compacto de estado recurrente. Este diseño mejora considerablemente la eficiencia tanto en el entrenamiento como en la inferencia a través del algoritmo de entrenamiento eficiente en hardware de GLA y el tamaño reducido del estado. Además, conservar la operación softmax es particularmente beneficioso en entornos de "ajuste fino de Transformers preentrenados a RNNs" (T2R), reduciendo la necesidad de un extenso entrenamiento desde cero. Experimentos extensos confirman el rendimiento superior de GSA en escenarios que requieren recuperación contextual y en entornos T2R.

English

Linear attention Transformers and their gated variants, celebrated for enabling parallel training and efficient recurrent inference, still fall short in recall-intensive tasks compared to traditional Transformers and demand significant resources for training from scratch. This paper introduces Gated Slot Attention (GSA), which enhances Attention with Bounded-memory-Control (ABC) by incorporating a gating mechanism inspired by Gated Linear Attention (GLA). Essentially, GSA comprises a two-layer GLA linked via softmax, utilizing context-aware memory reading and adaptive forgetting to improve memory capacity while maintaining compact recurrent state size. This design greatly enhances both training and inference efficiency through GLA's hardware-efficient training algorithm and reduced state size. Additionally, retaining the softmax operation is particularly beneficial in "finetuning pretrained Transformers to RNNs" (T2R) settings, reducing the need for extensive training from scratch. Extensive experiments confirm GSA's superior performance in scenarios requiring in-context recall and in T2R settings.