La hipótesis de separación de la predicción del estado

Resumen

Los Transformers utilizan el mismo flujo de cómputo hacia adelante tanto para predecir el siguiente token como para almacenar información de estado útil para predicciones futuras de tokens. Formulamos la hipótesis de separación entre estado y predicción: desvincular ambos roles produce un mejor rendimiento en modelado de lenguaje. Diseñamos una variante de Transformer que emplea dos flujos de cómputo para separar ambas funciones, y realizamos experimentos de preentrenamiento en diversas escalas. Nuestros experimentos muestran que la separación entre estado y predicción ofrece consistentemente mejores eficiencias en datos y cómputo, mejorando la pérdida de validación y superando a los Transformers estándar en 2 a 3 puntos porcentuales en promedio en tareas posteriores. También realizamos un análisis empírico exhaustivo que descarta posibles factores de confusión y demuestra la diferencia fundamental en los gradientes que implica nuestro diseño.

English

Transformers use the same forward computation stream to both predict the next token and store useful state for future token predictions. We formulate the state-prediction separation hypothesis: disentangling the two roles yields better language modeling performance. We design a Transformer variant that uses two computation streams to separate the two functions, and conduct pretraining experiments across various scales. Our experiments show that state-prediction separation consistently offers better data and compute efficiencies, improving validation loss and outperforming standard Transformers by 2--3 percentage points on average on downstream tasks. We also conduct extensive empirical analysis that rules out potential confounders and demonstrates the fundamental difference in the gradients our design entails.