UCSD y Together AI Research presentan Parcae: una arquitectura estable para modelos de lenguaje en bucle que logra la calidad de un transformador dos veces más grande

La receta dominante para construir mejores modelos de lenguaje no ha cambiado mucho desde la era Chinchilla: gastar más FLOP, agregar más parámetros, entrenar con más tokens. Pero a medida que las implementaciones de inferencia consumen una proporción cada vez mayor de computación y las implementaciones de modelos avanzan hacia el borde, los investigadores se hacen cada vez más una pregunta más difícil: ¿se puede escalar la calidad sin escalar el uso de memoria?

Un equipo de investigadores de UC San Diego y Together AI han presentado Parcae, una arquitectura de transformador en bucle estable que supera a los modelos en bucle anteriores y supera las líneas base de transformador de profundidad fija en cada escala probada, todo ello utilizando el mismo recuento de parámetros y el mismo presupuesto de datos de entrenamiento.

https://arxiv.org/pdf/2604.12946

¿Qué es un modelo de lenguaje en bucle?

En un Transformer estándar, las activaciones fluyen a través de una pila fija de capas exactamente una vez. En cambio, una arquitectura en bucle enruta las activaciones a través de un bloque de capas T veces en un bucle, multiplicando el cálculo efectivo sin agregar parámetros. Piense en ello como ejecutar el mismo grupo de bloques transformadores repetidamente en lugar de construir un modelo más alto.

Parcae utiliza específicamente un diseño de bucle medio, dividiendo la arquitectura en tres bloques funcionales: un preludio (P) que incorpora la secuencia de entrada en un estado latente e; un bloque recurrente (R) que actualiza iterativamente un estado oculto ht para bucles T, con e inyectado en cada iteración para mantener la influencia de la entrada; y una coda (C) que procesa el hT final para producir el resultado. Esta estructura mantiene el modelo compacto en la memoria, una propiedad valiosa para la implementación en el dispositivo, al tiempo que permite una cantidad significativamente mayor de cálculo por paso directo.

Los trabajos anteriores sobre transformadores en bucle, incluidos los modelos de profundidad recurrentes (RDM), fueron prometedores en un principio, pero fueron bastante difíciles de entrenar. Sufrieron una explosión de estado residual (donde el vector de estado oculto crece incontrolablemente a lo largo de las iteraciones del bucle) y frecuentes picos de pérdidas. Se requirió un ajuste sensible de hiperparámetros solo para lograr la convergencia.

La causa fundamental: un sistema residual sin restricciones

El equipo de investigación detrás de la idea clave de Parcae es reformular el paso hacia adelante del modelo en bucle como un sistema dinámico no lineal variable en el tiempo sobre el flujo residual:

ht+1 = Ā ht + B̄ e + R̄(ht, e),

Aquí, Ā controla el equilibrio entre los estados residuales anteriores y actuales, B̄ inyecta la señal de entrada y R̄ es la contribución no lineal de los bloques del transformador (atención y MLP). Al eliminar R̄ se obtiene un sistema lineal discreto invariante en el tiempo (LTI), y la teoría de control clásica proporciona inmediatamente la condición de estabilidad: el sistema es estable cuando la norma espectral ρ(Ā) < 1, marginalmente estable cuando ρ(Ā) = 1 e inestable cuando ρ(Ā) > 1.

El examen de los métodos anteriores bajo este marco revela precisamente el problema. Conjuntos de inyección de entradas basados ​​en sumas Ā = I (la matriz identidad), lo que significa ρ(Ā) = 1, marginalmente estable. El enfoque de concatenación con proyección utilizado por los RDM deja a Ā completamente libre, lo que hace que ρ(Ā) sea potencialmente mucho mayor que 1: inestable. Las curvas de entrenamiento empírico confirman esto directamente: las ejecuciones de entrenamiento divergentes aprenden ρ(Ā) ≥ 1, mientras que las pocas ejecuciones convergentes mantienen ρ(Ā) < 1.

Cómo Parcae impone la estabilidad mediante el diseño

En lugar de parametrizar Ā directamente, Parcae trabaja en forma continua y discretiza usando retención de orden cero (ZOH) y esquemas de Euler (tomando prestada una técnica estándar de modelos de espacio de estados como Mamba y S4) con un tamaño de paso aprendido Δ ∈ ℝdh, dando Ā = exp(ΔA) y B̄ = ΔB. Para garantizar ρ(Ā) < 1, la matriz continua A está restringida como una matriz diagonal negativa: A := Diag(−exp(logA)), donde logA ∈ ℝdh es un vector que se puede aprender. Debido a que las entradas diagonales siempre son negativas antes de la exponenciación, la restricción de la norma espectral se satisface en todo momento mediante la construcción.

Resultados: modelos con mejor rendimiento que el doble de tamaño

Frente a los RDM con parámetros y datos combinados entrenados en el conjunto de datos de Huginn, Parcae reduce la perplejidad de la validación hasta en un 6,3%, una cifra que alcanza su punto máximo en una escala de 350M (mejorando de 10,76 a 10,09 PPL) frente a una ganancia del 4,5% en una escala de 100M (14,23 a 13,59 PPL). La perplejidad de WikiText mejora hasta un 9,1% a una escala de 350M. La precisión promedio del punto de referencia de disparo cero en sentido descendente mejora hasta en 1,8 puntos.

Frente a las líneas de base estándar de Transformer de profundidad fija entrenadas con una configuración inspirada en nanochat en FineWeb-Edu, Parcae supera en todas las escalas. Con 1,3 mil millones de parámetros entrenados en tokens de 104 mil millones, Parcae supera al Transformer con parámetros coincidentes por 2,99 puntos en Core y 1,18 puntos en Core-Extended. El modelo 770M Parcae (25.07 Core) alcanza una calidad comparable a la del 1.3B Transformer (25.45 Core), aproximadamente la mitad de los parámetros para una capacidad equivalente. El equipo de investigación cuantifica la eficiencia de los parámetros de Parcae alcanzando hasta el 87,5% de la calidad de un transformador dos veces su tamaño, medida en comparación con la brecha de calidad con el siguiente modelo más grande.

Las primeras leyes de escala para bucles

La segunda contribución importante de esta investigación es el establecimiento de las primeras leyes de escala predecibles para el bucle de capas. Utilizando experimentos isoFLOP a escalas de 140M y 370M, el equipo de investigación muestra que el entrenamiento de cálculo óptimo aumenta la recurrencia media µrec y los tokens de entrenamiento D en conjunto, siguiendo leyes de potencia con exponentes consistentes en ambas escalas: µrec óptimo escala como C0.40 y los tokens óptimos escalan como C0.78, donde C es el presupuesto FLOP de entrenamiento.

Cuando los modelos Parcae en bucle entrenados en su µrec óptimo se comparan con modelos Parcae de profundidad fija (µrec = 1) bajo FLOP y presupuestos de parámetros idénticos, el bucle logra una pérdida de validación estrictamente menor, lo que se traduce en puntuaciones Core de 1,2 a 2,0 puntos más altas dependiendo del presupuesto FLOP. El bucle es un eje genuinamente ortogonal para escalar la computación, no un almuerzo gratis por compartir el peso.

En el momento de la prueba, el aumento del recuento de bucles T más allá de la profundidad del entrenamiento sigue una caída exponencial de saturación: L(T) = L∞ + Z·e−z·T, donde L∞ es un piso irreducible determinado por la profundidad del entrenamiento. Se estabiliza cerca de µrec (la recurrencia media utilizada durante el entrenamiento), lo que significa que la profundidad del entrenamiento establece un límite estricto en la escala en el tiempo de prueba. Estas dinámicas se unifican en una única ley paramétrica que predice la pérdida del modelo retenido dentro de un error promedio de 0,85 a 1,31%.

Conclusiones clave

Los transformadores en bucle ahora se pueden entrenar de manera confiable a escala: Parcae es una arquitectura en bucle para resolver los problemas de explosión de estado residual y picos de pérdida que han afectado a los modelos en bucle anteriores, logrando un entrenamiento estable en una amplia gama de tasas de aprendizaje donde los enfoques anteriores divergían. Un modelo Parcae 770M iguala la calidad de un transformador estándar 1.3B: al reutilizar las mismas capas en múltiples iteraciones de bucle en lugar de agregar más parámetros, Parcae ofrece una capacidad descendente equivalente con aproximadamente la mitad del espacio de memoria. El bucle es un tercer eje ortogonal para escalar la computación, junto con los parámetros y los datos: bajo un FLOP fijo y un presupuesto de parámetros, el entrenamiento óptimo para la computación requiere una recurrencia media creciente y tokens de entrenamiento en conjunto siguiendo leyes de potencia predecibles, lo que brinda a los profesionales de la IA una nueva palanca para mejorar la calidad sin comprar más hardware. Los bucles en el tiempo de prueba tienen un techo rígido establecido por la profundidad del entrenamiento: Parcae puede usar iteraciones de bucle adicionales en la inferencia para escalar el cálculo, pero se estabiliza cerca de la recurrencia media utilizada durante el entrenamiento. No se puede recorrer infinitamente el camino hacia un mejor rendimiento sin entrenar primero el modelo con recurrencias más profundas.

Consulte el papel, los pesos de los modelos y los detalles técnicos. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 130.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.

¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros