Presentan Differential Transformer V2 con mejoras en eficiencia y estabilidad para LLMs

Se presenta el Differential Transformer V2 (DIFF V2), una versión mejorada del Differential Transformer (DIFF V1). Este modelo se centra en la eficiencia de inferencia, la estabilidad de entrenamiento en LLMs a nivel de producción y la elegancia arquitectónica.

Las mejoras clave incluyen:

1. Inferencia más rápida y sin necesidad de núcleos de atención personalizados.
   DIFF V2 logra igualar la velocidad de decodificación del Transformer base y utiliza FlashAttention sin núcleos personalizados.
2. Mejor estabilidad en el entrenamiento.
   Se eliminó el RMSNorm por cabeza, lo que reduce la inestabilidad en etapas avanzadas del preentrenamiento de LLM.
3. Parametrización e inicialización más simples.
   Se reemplaza el parámetro global $\lambda lambda$ por uno específico de cada token y cabeza, lo que simplifica la inicialización.

Los experimentos de preentrenamiento con LLMs a gran escala muestran una notable reducción en la pérdida de modelado del lenguaje en comparación con el Transformer. Además, se observan menores picos de pérdida y gradiente durante el entrenamiento, especialmente con tasas de aprendizaje altas. Se espera evaluar el rendimiento en tareas de contexto prolongado y la eficiencia de aprendizaje en etapas posteriores.

Código

DIFF V2 se compara con DIFF V1 en los siguientes aspectos:

def DiffAttnV2(q, k, v, lam):
    attn = flash_attn_func(q, k, v)
    attn1, attn2 = (attn[:, 0::2], attn[:, 1::2])
    lam_val = sigmoid(lam)
    attn = attn1 - lam_val * attn2
    return attn

El código completo está disponible en: unilm/Diff-Transformer/Diff-Transformer-V2 en GitHub.

Observaciones Experimentales

Los experimentos de preentrenamiento realizados en LLMs de producción muestran una pérdida de modelado del lenguaje notablemente más baja en comparación con el Transformer, con una diferencia de 0.02 a 0.03 con 1T de tokens de entrenamiento. Se ha observado una reducción de picos de pérdida y gradiente, especialmente con tasas de aprendizaje grandes.

Las próximas evaluaciones se centrarán en la eficiencia de aprendizaje y el rendimiento en tareas de contexto prolongado.