ALiBi Positional Bias/es
| Article | |
|---|---|
| Topic area | transformers |
| Prerequisites | Transformer, Self-Attention, Positional Encoding |
Resumen
La Atención con Sesgos Lineales (ALiBi) es un método para inyectar información posicional en la atención del Transformer añadiendo un sesgo fijo y dependiente de la distancia a las puntuaciones de atención, en lugar de incrustar las posiciones en las representaciones de los tokens. Introducido por Press, Smith y Lewis en 2021, ALiBi reemplaza las codificaciones posicionales sinusoidales y las incrustaciones posicionales aprendidas por una matriz de sesgo estatica por cabeza cuyas entradas crecen linealmente con la distancia entre las posiciones de la consulta y la clave. Las cabezas se penalizan por atender a tokens lejanos a distintas tasas, controladas por una pendiente fija por cabeza.
La propiedad distintiva de ALiBi es su capacidad de extrapolar a longitudes de secuencia muy superiores a las vistas durante el entrenamiento. Las codificaciones posicionales estandar se degradan bruscamente al evaluarse con entradas mas largas que el contexto de entrenamiento, pero un modelo entrenado con longitud 1024 con ALiBi puede evaluarse con longitud 2048, 4096 o mayor con solo una degradacion modesta de la perplejidad, sin reentrenamiento ni interpolacion alguna. Esta propiedad hizo de ALiBi un caballo de batalla temprano para el modelado de lenguaje de contexto largo y el metodo posicional preferido en BLOOM, MPT y Replit Code.
Motivacion
El Transformer original utiliza codificaciones posicionales sinusoidales que se suman a las incrustaciones de los tokens de entrada. Las incrustaciones posicionales absolutas aprendidas (como en BERT y GPT-2) reemplazan las sinusoides fijas por vectores entrenables. Ambos enfoques comparten una limitacion fundamental: las posiciones mas alla del rango de entrenamiento o quedan fuera de la tabla de incrustaciones (aprendidas) o caen en partes de la sinusoide que el modelo nunca tuvo que interpretar (sinusoidales). Empiricamente, la perplejidad en posiciones fuera de rango se dispara.
Press, Smith y Lewis observaron que la informacion posicional de un Transformer entra en la red en ultima instancia a traves de las puntuaciones de atencion. En lugar de canalizar la posicion mediante las incrustaciones de entrada, ¿por que no codificarla directamente donde se consume? ALiBi elimina por completo las incrustaciones posicionales y, en su lugar, sesga los logits de atencion mediante una funcion de la distancia relativa entre la consulta y la clave. Como el sesgo es una funcion lineal simple de la distancia sin parametros aprendidos, evaluarlo a una distancia nunca vista esta bien definido y se comporta de forma suave.
Formulacion
La atencion por producto escalado estandar calcula, para consultas $ Q $, claves $ K $ y valores $ V $,
$ {\displaystyle \operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V.} $
ALiBi modifica esto añadiendo una matriz de sesgo $ B^{(h)} $ por cabeza dentro del softmax, antes de la normalizacion por filas,
$ {\displaystyle \operatorname{Attention}^{(h)}(Q, K, V) = \operatorname{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}} + B^{(h)}\right) V.} $
En el contexto del modelado de lenguaje causal, el sesgo para la cabeza $ h $ en la posicion de consulta $ i $ y la posicion de clave $ j $ con $ j \le i $ es
$ {\displaystyle B^{(h)}_{ij} = -m_h \cdot (i - j),} $
donde $ m_h > 0 $ es una pendiente fija (no aprendida) por cabeza y $ (i - j) \ge 0 $ es la distancia causal. Las posiciones $ j > i $ se enmascaran a $ -\infty $ como de costumbre. Para un modelo bidireccional, el sesgo es simetrico: $ B^{(h)}_{ij} = -m_h \cdot |i - j| $.
Las caracteristicas clave son: no se añade ninguna incrustacion posicional a las entradas; el sesgo depende solo de la distancia $ i - j $, no de las posiciones absolutas; y las pendientes $ m_h $ no se aprenden. Cada cabeza implementa, por tanto, una preferencia con decaimiento exponencial por los tokens cercanos, con la tasa de decaimiento controlada por su pendiente.
Cronograma de pendientes
ALiBi asigna las pendientes geometricamente entre las cabezas. Para una capa de atencion con $ n $ cabezas, las pendientes son
$ {\displaystyle m_h = 2^{-8 h / n}, \quad h = 1, 2, \ldots, n.} $
Para $ n = 8 $ cabezas, esto da pendientes $ 1/2, 1/4, 1/8, \ldots, 1/256 $. Para $ n = 16 $, las pendientes interpolan el mismo rango geometricamente, siendo cada pendiente la media geometrica de las pendientes superior e inferior de sus vecinas en el cronograma de 8 cabezas. Las cabezas con pendientes mayores atienden casi exclusivamente a tokens recientes; las cabezas con pendientes menores tienen un sesgo casi plano y pueden atender a contexto distante. El cronograma se eligio empiricamente por los autores y se mantuvo fijo a lo largo de los tamaños de modelo; pese a su sencillez, no se ha mejorado sustancialmente.
Es crucial que las pendientes sean constantes, no parametros. No se actualizan mediante Backpropagation, no aparecen en el estado del optimizador y añaden cero parametros al modelo. Esto hace que ALiBi sea mas barato que las incrustaciones posicionales aprendidas o el Rotary Position Embedding en cuanto a numero de parametros.
Extrapolacion de longitud
El resultado destacado del articulo original de ALiBi es la extrapolacion de longitud. Un Transformer de 1.3B parametros entrenado con contextos de 1024 tokens y evaluado con contextos de 3072 tokens conserva una perplejidad dentro de un margen pequeño con respecto a la que habria obtenido si se hubiera entrenado a 3072 tokens. Las codificaciones sinusoidales y las incrustaciones posicionales aprendidas, evaluadas con la misma longitud fuera de distribucion, ven explotar la perplejidad un orden de magnitud o mas.
El mecanismo es intuitivo. El sesgo $ -m_h (i - j) $ es una funcion lineal de la distancia y esta bien definido para cualquier $ i - j $ no negativo, incluidos valores mayores que cualquiera visto en el entrenamiento. El modelo nunca tiene que interpretar una entrada posicional que no haya encontrado; solo tiene que aplicar softmax a un sesgo lineal cuya forma es identica a la que siempre ha visto. La distribucion de atencion se concentra naturalmente en los tokens cercanos a larga distancia, en proporcion a la pendiente de cada cabeza, lo cual el modelo aprende a explotar durante el entrenamiento.
En la practica, la extrapolacion no es gratuita. La calidad continua degradandose gradualmente con la longitud, y las evaluaciones muy largas (mas de varias veces la longitud de entrenamiento) pierden informacion sobre tokens distantes porque los sesgos de todas las cabezas se vuelven grandes y negativos. Aun asi, la degradacion gradual hace de ALiBi una opcion solida cuando la longitud de contexto en entrenamiento es restringida pero el contexto en inferencia puede crecer.
Comparacion con otros metodos posicionales
Las codificaciones posicionales sinusoidales suman una funcion fija de la posicion a las incrustaciones de entrada. Son absolutas (cada posicion tiene una codificacion distinta), libres de parametros y, en principio, definidas para cualquier posicion; en la practica, no extrapolan porque la red entrenada no ha aprendido a interpretar las componentes de alta frecuencia mas alla del rango de entrenamiento.
Las incrustaciones posicionales absolutas aprendidas asignan un vector entrenable a cada posicion. No pueden extrapolar en absoluto; las posiciones mas alla de la tabla no existen.
Las codificaciones posicionales relativas (sesgo de T5, Shaw et al.) añaden un escalar aprendido como sesgo a las puntuaciones de atencion segun la distancia relativa cubeteada entre consulta y clave. ALiBi esta estrechamente relacionado, pero usa un sesgo lineal fijo en lugar de uno cubeteado y aprendido, lo que elimina parametros y permite extrapolar a distancias nunca vistas.
El embedding posicional rotatorio (RoPE) rota pares de dimensiones en $ Q $ y $ K $ por un angulo proporcional a la posicion. RoPE es ahora mas popular que ALiBi en los grandes modelos de lenguaje modernos (LLaMA, Mistral, Qwen, Gemma) porque tiende a dar una calidad ligeramente superior dentro de la distribucion y soporta extension de contexto basada en interpolacion mediante escalado consciente del NTK o YaRN. RoPE no extrapola tan limpiamente como ALiBi de fabrica, pero combinado con trucos de interpolacion de posiciones ha superado a ALiBi en el ajuste fino de contexto largo.
Variantes y extensiones
Se han propuesto varios refinamientos.
El ALiBi simetrico aplica $ B_{ij} = -m_h |i - j| $ para codificadores bidireccionales como los modelos al estilo BERT. El articulo original se centra en modelos causales, pero la variante simetrica se ha utilizado en configuraciones solo de codificador y de codificador-decodificador.
Las pendientes aprendibles hacen que $ m_h $ sea entrenable en lugar de fijo. Esto suele no mejorar la calidad y destruye parcialmente la garantia de extrapolacion, ya que las pendientes aprendidas pueden derivar en direcciones que dependen de la distribucion de entrenamiento.
Sandwich y KERPLE reemplazan el sesgo lineal por otras funciones monotonas de la distancia (logaritmicas, kernelizadas) que conservan la propiedad de extrapolacion al tiempo que a veces ofrecen una mejor perplejidad dentro de la distribucion. Estos metodos son conceptualmente descendientes de ALiBi y encajan en la misma posicion arquitectonica.
El ALiBi NTK Dinamico ajusta el cronograma de pendientes en tiempo de inferencia para extender aun mas el contexto efectivo, de forma analoga al escalado consciente del NTK para RoPE.
Notas de implementacion
ALiBi no requiere cambios en la canalizacion de entrada, ninguna capa de incrustacion posicional ni modificacion de $ Q $ o $ K $. Las implementaciones suelen precalcular la matriz de sesgo $ B^{(h)} $ una vez por longitud de secuencia y sumarla a los logits de atencion. En FlashAttention y kernels fusionados similares, el sesgo se calcula al vuelo a partir de los indices de fila y columna, usando un pequeño bufer de pendientes por cabeza de tamaño $ n $.
Para inferencia con cache de claves-valores, el sesgo depende solo de las posiciones absolutas de las claves cacheadas en relacion con la consulta actual, que se conocen. La inferencia en streaming con longitud de contexto creciente, por tanto, no necesita un manejo especial mas alla de extender el sesgo a las nuevas posiciones.
ALiBi es compatible con la Layer Normalization, las arquitecturas de decodificador y codificador, y el enmascaramiento de atencion estandar. Como no se mezcla informacion posicional en el flujo residual, la tabla de incrustaciones de entrada tambien es ligeramente menor que en los modelos con incrustaciones posicionales aprendidas.
Limitaciones
ALiBi tiene varias debilidades conocidas. En primer lugar, aunque la extrapolacion es gradual, la calidad dentro de la distribucion en pruebas donde importa la precision posicional (p. ej., recuperacion de spans con coincidencia exacta, tareas a nivel de caracter, ciertos casos de complecion de codigo) suele ser un poco inferior a la de un RoPE bien afinado. En segundo lugar, el cronograma de pendientes es fijo y no se adapta por tarea ni por capa; algunas cabezas terminan siendo efectivamente inutiles porque sus pendientes son demasiado pronunciadas para atender mas alla del entorno inmediato. En tercer lugar, el sesgo lineal con la distancia es un fuerte sesgo inductivo hacia la localidad; las tareas que requieren atender a un unico token distante (p. ej., recuperar una clave del comienzo de un contexto largo) se vuelven mas dificiles a medida que crece el contexto, ya que el sesgo suprime activamente las claves lejanas.
Estas concesiones explican por que ALiBi ha sido parcialmente desplazado por RoPE combinado con metodos de interpolacion de posiciones en los grandes modelos de lenguaje recientes, aunque sigue siendo una linea base limpia y libre de parametros para la extrapolacion de longitud y todavia se utiliza activamente en sistemas en produccion entrenados antes de 2023.
Referencias
- ↑ Press, O., Smith, N. A., and Lewis, M. "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation." International Conference on Learning Representations (ICLR), 2022. arXiv:2108.12409.
- ↑ Vaswani, A. et al. "Attention Is All You Need." Advances in Neural Information Processing Systems (NeurIPS), 2017. arXiv:1706.03762.
- ↑ Shaw, P., Uszkoreit, J., and Vaswani, A. "Self-Attention with Relative Position Representations." North American Chapter of the Association for Computational Linguistics (NAACL), 2018. arXiv:1803.02155.
- ↑ Su, J. et al. "RoFormer: Enhanced Transformer with Rotary Position Embedding." Neurocomputing, 2024. arXiv:2104.09864.
- ↑ Chi, T.-C., Fan, T.-H., Ramadge, P. J., and Rudnicky, A. I. "KERPLE: Kernelized Relative Positional Embedding for Length Extrapolation." Advances in Neural Information Processing Systems (NeurIPS), 2022. arXiv:2205.09921.
- ↑ Le Scao, T. et al. "BLOOM: A 176B-Parameter Open-Access Multilingual Language Model." 2022. arXiv:2211.05100.