Implementación de la dispersión con la función Top-K
La operación más fundamental que realiza Q-Sparse es aplicar la función de dispersión Top-K al tensor de entrada.
Específicamente, la arquitectura Transformer utiliza capas lineales nn.Linear (multiplicación de matrices) para proyecciones en las capas de atención y de alimentación hacia adelante, que se pueden representar como Y=X·W^T. (Donde X es el tensor de entrada, W representa sus pesos, e Y es el tensor de salida)
En Q-Sparse, para un tensor de activación de entrada X, primero se calcula su valor absoluto |X| y se ordena para encontrar los K elementos con el mayor valor absoluto.
Aquí, K es un hiperparámetro predefinido que determina el grado de dispersión.
Luego, Q-Sparse crea un tensor de máscara binaria M con la misma forma que X, estableciendo las posiciones correspondientes a los K elementos con mayor valor absoluto en |X| a 1, y el resto a 0.
A continuación, se realiza el producto Hadamard (multiplicación elemento a elemento) entre el tensor de entrada X y el tensor de máscara M para obtener el tensor disperso X_sparse.
Durante la propagación hacia adelante, el tensor disperso X_sparse reemplaza al tensor de entrada original X en los cálculos subsiguientes (como la multiplicación de matrices).
Dado que la mayoría de los elementos en X_sparse se han establecido a cero, esto reduce significativamente la cantidad de cálculos y los requisitos de ancho de banda de memoria.
Durante la retropropagación, Q-Sparse utiliza el estimador de paso directo (Straight-Through Estimator, STE) para calcular el gradiente de la función Top-K.
En el entrenamiento tradicional, generalmente se necesita calcular el gradiente de la función de pérdida con respecto a los parámetros de la red y usar el descenso de gradiente para actualizar los parámetros y minimizar la pérdida.
Sin embargo, cuando existen operaciones no diferenciables como la cuantización o Top-K en la red, el cálculo del gradiente se vuelve problemático, ya que el gradiente de la salida con respecto a la entrada es cero en la mayoría de los puntos para estas operaciones, lo que impide una propagación efectiva del gradiente.
El STE evita el problema de desvanecimiento del gradiente pasando el gradiente directamente al tensor antes de la dispersión.
En la retropropagación general, el gradiente de la función de pérdida L con respecto a x es ∂L/∂x=∂L/∂y⋅∂y/∂x, pero esto no se puede calcular directamente debido a la no diferenciabilidad.
La solución del STE es calcular solo el gradiente de la función de pérdida con respecto al tensor disperso y, y luego copiarlo directamente al tensor original x, es decir, usar ∂L/∂y directamente como una estimación de ∂L/∂x.
Para las capas de alimentación hacia adelante, Q-Sparse utiliza la función ReLU cuadrada en lugar de la función de activación ReLU regular, donde la operación cuadrática puede aumentar aún más la dispersión de la activación (⊙ representa el producto Hadamard).
Además, para adaptarse a los modelos cuantizados, Q-Sparse cuantiza el tensor de entrada antes de aplicar la dispersión Top-K para asegurar que la operación de dispersión sea compatible con la representación cuantizada, expresada por la siguiente función:
Donde ε es una pequeña constante para evitar divisiones por cero.
En particular, para pesos cuantizados de 1 bit, Q-Sparse utiliza la siguiente función de cuantización, donde α es el valor absoluto promedio del tensor de pesos W.
60% de parámetros de activación logran el mismo efecto
Los experimentos comparativos muestran que Q-Sparse supera significativamente a los métodos ReLU anteriores tanto en tasa de dispersión como en rendimiento del modelo.
Para los efectos específicos de Q-Sparse, los autores evaluaron su rendimiento en tareas de entrenamiento desde cero, continuación del entrenamiento y ajuste fino.
Los experimentos de entrenamiento desde cero utilizaron el modelo Llama, y los resultados mostraron que Q-Sparse con un 70% top-K (es decir, una tasa de dispersión general del 40%) podía alcanzar una pérdida de entrenamiento comparable a la línea base densa en modelos de 700M y 7B.
El propósito de la continuación del entrenamiento es dispersar modelos densos, y el objeto de experimento aquí fue Mistral-7B.
Como resultado, con parámetros de activación de 2.9B y 3.8B, las puntuaciones del modelo en conjuntos de datos como ARC y MMLU no mostraron una disminución significativa.
En los experimentos de ajuste fino, para los modelos Qwen-7B y Mistral-7B, Q-Sparse mostró resultados similares a la continuación del entrenamiento, logrando un rendimiento muy cercano al de los modelos densos con alrededor del 60% de los parámetros de activación.
Estos resultados implican que, con el mismo rendimiento, los modelos de activación dispersa pueden reducir significativamente los parámetros de activación durante la inferencia en comparación con los modelos densos, reduciendo así la cantidad de FLOPS consumidos.
Para los modelos cuantizados, el equipo aplicó Q-Sparse a su modelo BitNet b1.58 desarrollado internamente y lo entrenó y evaluó en múltiples conjuntos de datos.
Se puede observar que, en escalas de 700M y 7B, la velocidad de convergencia y el valor final de la función de pérdida de los modelos cuantizados que utilizan Q-Sparse son comparables a los de los modelos cuantizados sin Q-Sparse (BitNet b1.58).
Esto indica que Q-Sparse se puede integrar perfectamente en los modelos cuantizados sin afectar significativamente el entrenamiento y la convergencia del modelo.
Basándose en esto, los autores creen que combinar Q-Sparse con técnicas de cuantización puede mejorar aún más la eficiencia de los grandes modelos de lenguaje durante la fase de inferencia.
Descubrimiento de una nueva "Ley de Escalado" para la optimización de inferencia
Además de evaluar el rendimiento de estos modelos con activación dispersa, los autores también exploraron la relación entre el rendimiento del modelo, la escala y la tasa de dispersión, y realizaron algunos nuevos descubrimientos.
Ley de escalado de rendimiento para modelos de activación dispersa: Los autores descubrieron que, similar a los modelos densos, el rendimiento de los modelos de activación dispersa también sigue una relación de ley de potencia.
Específicamente, dado una tasa de dispersión S, el valor de la función de pérdida L(N,S) en la convergencia del modelo se puede aproximar mediante la siguiente fórmula: