Differential Transformer: Transformer que cancela el ruido de atención

 (arxiv.org)
1 puntos por GN⁺ 2024-10-09 | 1 comentarios | Compartir por WhatsApp
  • el Transformer solo con decoder se ha convertido en la arquitectura estándar de los LLM, pero en contextos largos su rendimiento se ve afectado por el ruido de atención, al perder información clave y repartir atención a tokens irrelevantes
  • la atención diferencial divide query y key en dos grupos, crea dos mapas de atención softmax y luego resta el segundo, escalado con una λ aprendible, para calcular el score final de atención
  • DIFF Transformer mostró mejores resultados que Transformer en experimentos donde se aumentó el tamaño del modelo, los tokens de entrenamiento y la longitud de contexto, y necesita solo alrededor de 65% del tamaño de modelo o de los tokens de entrenamiento para lograr un rendimiento similar
  • muestra ventajas en modelado de contexto largo, recuperación de información clave, mitigación de alucinaciones, in-context learning, razonamiento matemático y reducción de activation outliers; en QA y resumen disminuye el impacto del contexto irrelevante
  • como mantiene el mismo diseño general que Transformer y solo reemplaza la atención softmax, puede reutilizar FlashAttention manteniendo equivalentes el número de parámetros y el costo computacional

El problema del ruido de atención en Transformer

  • el Transformer solo con decoder es la arquitectura estándar de facto para los LLM, y en su núcleo está el mecanismo de atención, que pondera con softmax la importancia de los tokens dentro de una secuencia
  • los LLM tienen dificultades para encontrar con precisión la información clave en el contexto, y mientras más contexto irrelevante haya, más fácil es que las pistas correctas queden enterradas
  • en ejemplos donde hay que encontrar una respuesta insertada en medio de un montón de documentos, Transformer tiende a asignar solo un pequeño attention score a la respuesta correcta y a repartir scores excesivos al contexto irrelevante
  • esos attention scores no despreciables asignados al contexto irrelevante actúan como ruido de atención
  • el ejemplo de Multi-Needle Retrieval de la Figura 1 muestra una precisión de 55% para Transformer y 85% para Differential Transformer

Cómo funciona la atención diferencial

  • DIFF Transformer es una arquitectura base para sequence modeling y LLM que mantiene el macro layout del Transformer existente y reemplaza la atención softmax normal por atención diferencial
  • a partir de la entrada X, proyecta query, key y value, pero divide query y key en dos grupos: Q1, Q2, K1, K2, mientras que value se mantiene como V
  • la salida de atención se calcula como la diferencia entre dos mapas de atención softmax
    • DiffAttn(X) = (softmax(Q1K1^T / √d) − λ softmax(Q2K2^T / √d))V
    • la estructura elimina el ruido común al restar el segundo mapa de atención del primero
  • λ es un escalar aprendible, y se reparametriza así para ajustar la dinámica de entrenamiento
    • λ = exp(λq1 · λk1) − exp(λq2 · λk2) + λinit
    • en los experimentos base se usa λinit = 0.8 − 0.6 × exp(−0.3 · (l − 1))
    • también se evaluó usar el mismo λinit para todas las capas, por ejemplo 0.8, y en las ablaciones mostró un rendimiento relativamente robusto ante diferencias en la estrategia de inicialización
  • este método es similar a la idea del amplificador diferencial, que elimina ruido de modo común a partir de la diferencia entre dos señales, y a los audífonos con cancelación de ruido
  • Naderi et al. demostraron que la atención diferencial hace más equilibrada la distribución espectral de la matriz de atención y resuelve eficazmente el colapso de rango

Estructura multi-head y arquitectura general

  • la atención diferencial multi-head usa distintas matrices de proyección para cada head, y dentro de una misma capa comparte el escalar λ entre heads
  • a la salida de cada head se le aplica RMSNorm de forma independiente, luego se multiplica por (1 − λinit), después los heads se concatenan en la dimensión de canales y finalmente pasan por la proyección de salida WO
  • la notación GroupNorm de la Figura 2 enfatiza que la normalización se aplica de manera independiente a cada head
    • la atención diferencial tiende a tener patrones más dispersos, por lo que la información estadística entre heads es más diversa
    • la normalización por head mejora las estadísticas del gradiente al normalizar cada head antes de la concatenación
  • toda la capa de DIFF Transformer se compone de dos módulos
    • MultiHead(LN(Xl)) + Xl
    • SwiGLU(LN(Yl)) + Yl
  • estructuralmente usa pre-RMSNorm y SwiGLU, siguiendo las mejoras de la familia LLaMA

Eficiencia y estabilidad de entrenamiento

  • la atención diferencial puede reutilizar directamente FlashAttention, lo que permite mejorar mucho la eficiencia del modelo
  • el número de heads se define como h = dmodel / 2d, donde d es igual a la dimensión de head del Transformer
  • esta configuración busca igualar con Transformer el número de parámetros y la complejidad computacional
  • después de la normalización por head se usa el multiplicador fijo (1 − λinit) para alinear el flujo de gradiente con Transformer
  • el Apéndice G muestra que el flujo total de gradiente se mantiene similar al de Transformer, lo que permite heredar hiperparámetros parecidos y asegurar estabilidad de entrenamiento

Resultados experimentales y efectos en la práctica

  • los experimentos de modelado de lenguaje escalan DIFF Transformer aumentando el número de parámetros, los tokens de entrenamiento y la longitud de contexto
  • en la scaling curve, DIFF Transformer necesita solo alrededor de 65% del tamaño de modelo o de los tokens de entrenamiento para alcanzar un rendimiento de modelado de lenguaje similar al de Transformer
  • también mostró mejor rendimiento que Transformer en tareas downstream, y en evaluaciones de secuencias largas aprovecha mejor el contexto a medida que este crece
  • en recuperación de información clave muestra un patrón de asignar mayores attention scores al span correcto y menores scores al contexto irrelevante
  • en QA y resumen de texto se distrae menos con contexto irrelevante, mitigando las alucinaciones
  • en in-context learning mejora la precisión y además es más robusto frente a la permutación del orden, considerada un problema crónico de robustez
  • la reducción de activation outliers abre nuevas oportunidades para la cuantización

Lecturas relacionadas

1 comentarios

 
GN⁺ 2024-10-09
Comentarios en Hacker News

  • Siento que aquí se está perdiendo la intuición clave. Entiendo el problema de que la atención softmax normal tiene dificultades para asignar una atención cercana a 0 a cosas irrelevantes, y también entiendo que, si hay una estructura de resta, se pueden crear pesos de atención exactamente o casi iguales a 0 sin valores de activación atípicos
    Aun así, esta estructura también parece que podría crear fácilmente pesos de atención negativos, y eso se ve equivalente a dar atención positiva a la negación del vector de valores. Intuitivamente, parece difícil mantener en equilibrio todo lo que no interesa cerca de 0. Aun así, Figure 1 muestra que funciona bien, así que no dudo de su potencial, pero me cuesta imaginar específicamente cómo la red logra hacerlo

    • Hay un error en la softmax y en la atención normales. Softmax debería ser exp()/1+∑exp()
      Lo clave es que se añade 1 al denominador. En el límite negativo, softmax puede volverse 0 en vez de algún épsilon. También se puede lograr el mismo efecto agregando un valor 0 adicional a x. La desventaja es que, para corregir esto, habría que volver a entrenar el modelo desde cero
    • En lugar de permitir pesos negativos, como no pasa por algo como ReLU, sí suena como si estorbara un poco al modelo. Pero este procesamiento podría ser un problema más fácil de lo que parece para el modelo
      Si miras la primera figura de pesos de atención, en la zona de ruido sí hay puntajes negativos. Aun así, la atención sobre esa zona de todos modos es muy pequeña. El segundo mapa de atención solo necesita predecir el ruido del primero, y como tiene acceso completo a la primera entrada, es una tarea que puede hacer con bastante precisión
      Volviendo a la analogía del artículo, los audífonos con cancelación de ruido pueden producir una señal de cancelación precisa porque el micrófono tiene acceso al sonido que llega al oído. De la misma manera, el segundo mapa de atención sabe qué entra en el primero, así que puede generar la señal de cancelación correspondiente. No es perfecto, pero los audífonos con cancelación de ruido tampoco lo son, y aun así llegar al 99% basta para mejorar el rendimiento
    • Intuitivamente, parece muy fácil que durante el entrenamiento el modelo optimice lambda a 0. Entonces, en esencia, se convierte en un Transformer normal con un mecanismo de poda de parámetros excesivamente complicado añadido
      La poda ya está bastante establecida en la literatura como una forma sorprendentemente buena de reducir la cantidad de parámetros, y puede bajar alrededor de un 40%. El modelo real probablemente no funcione exactamente así, pero tampoco sorprendería que al final terminara aproximándose a un Transformer normal
    • Los valores negativos pueden aumentar la capacidad de representación

  • Muy ingenioso. Me gustan este tipo de trabajos detallistas, y además el cambio es pequeño, así que parece algo que otros podrían aplicar fácilmente. Excelente
    Dicho eso, la última oración de la introducción de la sección "2 Differential Transformer" me preocupa un poco. Dicen que usan mejoras de artículos previos, pero por el contexto gramatical no queda claro si esas mejoras se aplicaron tanto al Transformer normal como al diff Transformer. Si no fue así, la comparación se vuelve confusa. La expresión "main difference" en la oración anterior me hizo prender las alarmas
    Claro, un investigador de buena fe podría saber esto y no sentir la necesidad de dejarlo explícito. Pero en parte de la investigación publicada en este campo, por más cuidadoso que seas nunca es demasiado

    • Sí. De verdad se ve muy bien. Hay mejoras de perplexity en general en tiempo de entrenamiento, por token de entrenamiento y por tamaño de modelo
      Me recuerda a las arquitecturas MoE, donde se elige el mejor modelo pequeño para encargarse de parte o de toda la tarea de inferencia. Me pregunto si MoE obtiene una ganancia parecida al forzar al Transformer a distinguir entre posibilidades alternativas
      En cualquier caso, si los números se mantienen, parece que se adoptará ampliamente. Como dije, prácticamente no se le ven desventajas y además parece fácil de reproducir
    • Los otros dos cambios que mencionan ya fueron adoptados ampliamente, y algunos de los modelos usados para comparar ya los incluyen. Parece que solo enumeraron los cambios respecto a la arquitectura Transformer original por completitud

  • Como con la mayoría de las cosas en este nuevo mundo del aprendizaje automático, resulta realmente confuso por qué esto funciona
    La analogía de los audífonos con cancelación de ruido ayuda, pero en ese caso está claro qué es la señal y qué es el ruido. Si aquí también lo supiéramos, entonces no queda claro por qué habría que hacer el trabajo de eliminar el ruido en primer lugar

    • Con un solo softmax no se puede predecir exactamente 0, solo números muy pequeños. Si hay muchos valores que sumar, esos valores pequeños terminan mezclando muchas cosas irrelevantes en la salida, contaminándola con el ruido del que habla el paper
      Peor aún, el gradiente de los valores de atención bajos se vuelve muy pequeño, así que deshacer esos errores requiere muchas actualizaciones de pesos. En cambio, al restar las salidas de dos softmax, el modelo puede predecir pesos exactamente iguales a 0 para algunos valores y mantener un flujo de gradiente razonable
      O sea, el modelo ya sabe qué es ruido, pero un solo softmax hace difícil excluirlo. Además, con un solo softmax se obliga a que la salida de todos los heads permanezca dentro de la envolvente convexa de los vectores de valor, mientras que en esta variante cada head puede elegir su propio lambda y mover el rango de salida fuera de la envolvente convexa definida de antemano por esos valores. Así que aumenta la capacidad de representación del modelo
    • Los audífonos con cancelación de ruido probablemente son una mala analogía aquí
      Un mejor ejemplo sería la señal diferencial usada en audio profesional y en muchos protocolos de señal digital como Ethernet, HDMI y USB. En vez de usar un solo cable con referencia a tierra, la señal se transmite como la diferencia entre dos cables. Ambos cables llevan la misma señal con polaridad opuesta y, como pasan uno junto al otro, el ruido externo se aplica por igual a ambos
      El voltaje cambiará, pero la diferencia de voltaje entre los dos cables se mantiene. En el receptor, al restar ambos voltajes, el ruido simplemente se cancela
    • En vez de buscar una analogía, basta con verlo como una nueva capacidad matemática añadida. Hace posible la atención negativa, permitiendo que la red diga en el cálculo de atención: "quiero restar la contribución de este token". Antes solo podía reducir cuánto sumar
      Una forma simple de hacer esto sería quitar el softmax o usar sigmoid, pero en la práctica parece que softmax funciona mejor
    • Una hipótesis de por qué esto funciona es que mitiga una desventaja de RoPE
      Dicho de forma simple, RoPE es una estrategia moderna que le da al modelo información sobre qué tan separadas están la consulta y la clave al hacer atención. Es la mejor estrategia disponible hoy, pero tiene el gran defecto de hacer que algunas conexiones entre tokens lejanos sean mucho más fuertes de lo deseado. Xpos (https://arxiv.org/pdf/2212.10554) también es un paper de Microsoft que aborda problemas de RoPE, y en la página 4, Figure 1, se puede ver una interpretación visual de la intensidad de atención con forma sinusoidal. Idealmente, uno querría que fuera suave
      Creo que una gran razón por la que Differential Transformer funciona especialmente bien en secuencias largas es que, incluso cuando q1 y q2 no coinciden con ningún token, la intensidad relativa de RoPE tiene el mismo valor, así que el ruido se cancela. Solo queda la coincidencia intencional, aunque al costo de debilitar en cierta medida el valor que RoPE aportaba originalmente
      Claro, esto es solo una hipótesis. Sería fácil verificarlo comparando experimentos donde ambos usen como baseline alibi attention (https://arxiv.org/pdf/2108.12409). alibi tiene otros compromisos que este método no mitiga, pero aun así es un resultado realmente interesante
    • Parte del trabajo previo aquí son las ladder networks y, de manera algo gestual, también las residual nets. Ambas pueden interpretarse como modelos entrenados para reducir el error de predicciones anteriores, en lugar de predecir directamente el resultado final
      La intuición de por qué funciona parece estar en que hace el paisaje del descenso por gradiente un poco más amigable y facilita aprender en pasos pequeños. Ahora la propia red se diseña explícitamente en torno a la idea de que al principio cometerá muchos errores en sus predicciones y mejorará con el tiempo

  • Si entendí bien "Differential attention takes the difference between two softmax attention functions to eliminate attention noise", esta arquitectura parece ser un intercambio donde se usa el doble de memoria de atención para obtener un modelo de mayor calidad, o una calidad similar con menos parámetros
    La parte que dice "6.8B-size DIFF Transformer achieves a validation loss comparable to 11B-size Transformer, requiring only 62.2% of parameters" me deja varias preguntas. Si solo usa el 60% de los parámetros, me pregunto si eso compensa el hecho de tener el doble de espacio de atención y termina con características de memoria parecidas a las de un Transformer tradicional, y si ese intercambio cambia de forma notable entre entrenamiento e inferencia

    • Entendí que los parámetros extra necesarios para el segundo mecanismo de atención también están incluidos en esos 6.8B parámetros. Es decir, no es una cantidad hipotética de parámetros que tendría un Transformer estándar, sino el total real del modelo. Así que el resultado es el doble de impresionante
      El paper dice: "We set the number of heads h = dmodel/2d, where d is equal to the head dimension of Transformer. So we can align the parameter counts and computational complexity." En otras palabras, compensan esto reduciendo a la mitad la cantidad de heads de atención por capa
    • Parece que reducen a la mitad la cantidad total de heads y duplican V y O para mitigar la memoria y el cómputo extra. No he revisado las cuentas exactas, pero salvo operaciones baratas como multiplicaciones por constantes y restas, parece que la cantidad de operaciones en coma flotante sería equivalente
    • El ahorro de RAM probablemente se compensa, pero el espacio necesario para almacenamiento sí se reduce, y también podría bajar el tiempo de arranque inicial dependiendo de la velocidad del almacenamiento y del tamaño del modelo. Así que podría estar bien para modelos modestos en dispositivos de consumo
    • El tamaño de la KV cache se duplica, y con contextos grandes esto puede convertirse en varios GB bastante rápido

  • Me da curiosidad qué historia hubo detrás de la fórmula "We empirically find that the setting λᵢₙᵢₜ = 0.8 − 0.6 × exp(−0.3 · (l − 1)) works well in practice"

    • 0.8 funciona bien, pero probemos con un valor inicial más bajo en las capas inferiores. Pongámosle algo como 0.2. Bien, ahora hace falta una fórmula que vaya y venga entre 0.2 y 0.8, acercándose lentamente a 0.8. Parece el tipo de cosa que salió de mover números durante 20 minutos hasta decidir que eso ya estaba bien
    • Muchísimas cosas se ajustan de forma óptima girando perillas analógicas o tanteando hasta que suenan bien
    • Al menos al inicio del entrenamiento, esta fórmula parece hacer que el término de atención negativo sea menor en las capas del frente, o sea, cuando l es pequeño, que en las capas posteriores. Suena razonable. Uno querría prestar un poco de atención a todo antes de concluir cuáles son las posiciones que realmente vale la pena mirar
      Aunque no parece que el autor discuta esta elección por separado en el paper

  • La parte clave que al principio no entendía era qué pasa si los dos grupos de atención aprenden lo mismo. Como las máscaras de atención se restan entre sí, si ambos producen valores parecidos, la atención total cae a 0 y la pérdida aumenta
    Así que la única forma de reducir la pérdida es aprender a prestar atención a cosas distintas. Una de las estrategias más simples que podrían aprender es, como sostiene el paper, que un grupo se enfoque en el contexto relevante y el otro en el contexto irrelevante. Así, un grupo aprende el ruido y el otro la señal. En la práctica no es tan tajante, pero como simplificación para entenderlo sirve

    • La parte interesante es que no se resta simplemente, sino que solo se resta una parte del segundo softmax
      Tiene sentido si piensas que, si las dos copias fueran idénticas, la salida del softmax también lo sería y la diferencia daría 0 en todas partes. Pero si restas una copia escalada, el proceso de normalizar la diferencia parece resaltar más los valores de la señal que los del ruido, haciendo que la señal destaque más que antes de la normalización
    • Me pregunto si habrá alguna analogía con nuestra propia experiencia de sorpresa y su utilidad para entender qué pasa cuando los dos grupos de atención aprenden lo mismo
      Sería como que una cabeza de atención sube el peso si le sorprende lo que aprendió la otra, y si ambas encuentran lo mismo, lo considera poco sorprendente y baja el peso
      Admito que la "sorpresa" ocupa una parte bastante grande de mi base de conocimientos[1][2][3]. Es una emoción subjetiva y una función adaptativa de la mente, uno de los sistemas adaptativos más complejos que conocemos
      [1] https://plus.maths.org/content/information-surprise
      [2] https://blakeelias.name/papers/Multi-Agent-Cooperation-Intri...
      [3] https://complexity.simplecast.com/episodes/81/transcript
    • Puede que haya una pequeña posibilidad de que ambos aprendan lo mismo, pero no parece lo bastante probable como para ser un problema importante
    • ¿No podría la función de pérdida penalizar que ambos aprendan lo mismo?

  • Me pregunto qué se pierde aquí. Seguro que hay algún trade-off
    También me pregunto si esto afecta la creatividad o la capacidad de interpolar entre conceptos. La alucinación y la creatividad parecen bastante relacionadas. Entiendo la alucinación como algo desalineado con el espacio de interpolación que los humanos sentimos como apropiado

    • No veo por qué la alucinación y la creatividad parecerían relacionadas. Yo lo veo simplemente como error de muestreo
      Claro, los errores a veces pueden inspirar, pero la creatividad es muchísimo más que cometer errores
      Estos modelos de lenguaje son predictores del siguiente token. Ese siguiente token se predice muestreando desde el espacio de probabilidades que produce el modelo. Ese proceso de muestreo puede ser no determinista
      Una alucinación es cuando ese muestreo produce tokens que forman una oración falsa o no intencionada. Incluso podría decirse que todo lo que produce el modelo es una alucinación, pero entrenamos al modelo para que emita espacios donde sea más probable que alucine lo que queremos. Si no, solo arrojaría ruido sin sentido
      "Alucinación" es una palabra realmente malísima para describir esto
    • Uno de los trade-offs es la velocidad y la memoria. Hay el doble de pesos Q y K en el bloque de atención, así que el throughput en su H100 bajó alrededor de un 10%. Aparece en la tabla 7 del apéndice A
    • No toda alucinación es creatividad. Si imaginas una aplicación RAG, el modelo tiene que apegarse a los documentos dados

  • Me pregunto qué parte del valor aquí viene de compensar el ruido posicional que introduce RoPE. Me gustaría ver una tabla que compare no solo los modelos con RoPE de este trabajo, sino también una versión con alibi y una línea base con alibi
    Aun así, es una mejora enorme, y felicidades a los investigadores

  • ¿Lo que está pasando aquí es que softmax no puede empujar valores hasta 0, pero al restar dos mapas softmax sí se puede producir una salida 0?

    • La pregunta siguiente es si la probabilidad de producir 0 no será extremadamente baja
    • O incluso valores negativos son posibles

  • Es un buen problema para resolver, pero creo que el enfoque está equivocado
    Para saber a qué se le prestó atención y conocer el contexto completo, esto debería hacerse de forma jerárquica. Si el vector diferencial se calcula a partir de la misma entrada que el vector de atención, no veo cómo puede saber correctamente cómo modificar el vector de atención

    • Al final, ¿no se ajusta todo en la dirección que indican las derivadas de backpropagation y en proporción a ese gradiente? En otras palabras, mientras el sistema de backprop funcione, parecería que no debería ser un problema decidir en qué dirección ajustar los pesos