Un transformer en bucle de 1B parámetros resultó 1,4 veces más barato en FLOP

Los transformers en bucle se han convertido este año en un tema de moda: casi cada día salen artículos. La idea es sencilla: en lugar de añadir capas nuevas, el modelo pasa varias veces por las mismas. Los parámetros son los mismos, pero "piensa" con más profundidad. Un nuevo trabajo de 149 páginas, firmado por investigadores entre los que están Mike Lewis y Artidoro Pagnoni, responde por fin a una pregunta práctica: cómo hacer el bucle exactamente para que compense.

Cuál es el problema del bucle

Imagine una cadena de producción con 24 máquinas. Puede comprar otras 24 y lograr un procesamiento más complejo, pero es caro en memoria y parámetros. O bien puede hacer pasar la pieza por las mismas máquinas una segunda vuelta. Se necesita la misma memoria, pero más cómputo. Este enfoque tiene dos puntos débiles. El primero: no está claro qué repetir, todas las capas o solo una parte. El segundo: al generar texto, cada nuevo token espera a que terminen todos los bucles del anterior, y paralelizar esa decodificación es difícil.

Qué descubrieron los autores

El equipo preentrenó decodificadores densos de hasta 1B de parámetros y hasta 32 bucles, y comparó dos modos. Full-Stack Looping repite todas las capas. Middle-Stack Looping repite solo las intermedias, y las primeras y las últimas quedan como están.

La conclusión depende de lo que le falte. Si el número de parámetros está estrictamente limitado, es mejor repetir todas las capas. Si lo limitado es el cómputo (FLOP), ganan las capas intermedias. Para un ingeniero, esto se convierte de inmediato en una regla para elegir la arquitectura.

La segunda parte del trabajo trata de la decodificación paralela. En el esquema Parallel Loop Transformers (PLT), los estados de los bucles se desplazan a los tokens vecinos, de modo que distintos bucles se calculan a la vez. Los autores notaron que con varios bucles la calidad se degrada, porque en una sola suma se mezclan el embedding del token actual y el estado del anterior. Su solución, Normalized Shifted PLT, normaliza por separado esas dos entradas mediante RMSNorm antes de mezclarlas. Además, añadieron adaptadores compartidos de bajo rango para los bucles (Loop Matrix), en esencia una corrección similar a LoRA, para que cada pasada difiera un poco de la anterior.

Las cifras son las siguientes. En el conjunto DCLM con 200M de parámetros, la versión de cuatro bucles reduce los bits por byte a ~1,060, frente a ~1,087 del PLT original y ~1,088 del modelo sin bucles. Con 1B de parámetros, el modelo de dos bucles logra ~0,86 frente a ~0,90 con 80 tokens por parámetro. La precisión en ocho tareas fue de ~47,5% frente a ~46,1% con 160 tokens por parámetro. Con entrenamiento prolongado, los modelos de dos bucles requieren hasta 1,4 veces menos FLOP que los modelos sin bucles.

Qué hacer con esto

Si no entrena modelos desde cero, ahora mismo no hay nada que hacer. Pero si planea un modelo compacto propio para un dispositivo o un servidor con un límite estricto de memoria, el bucle se convierte en una palanca real. Se necesita la misma memoria y la calidad aumenta. El esquema de elección es simple: si le limitan los parámetros, tome el bucle completo; si le limitan los FLOP, repita la parte central. Si importa la velocidad de generación, fíjese en el desplazamiento normalizado: es lo que hace que la decodificación paralela funcione.

Limitaciones

Hay tres salvedades importantes. Primero, los experimentos se detuvieron en 1B de parámetros y se desconoce si las tendencias se mantendrán en modelos de decenas de miles de millones. Segundo, los autores no midieron el tiempo real de generación. Ahorrar FLOP aún no significa que la respuesta llegue más rápido: en la práctica, la memoria y los kernels deciden mucho. Tercero, la mejora de precisión es pequeña, de alrededor de 1,4 puntos porcentuales, así que es más una optimización de ingeniería que un avance revolucionario.

En resumen: el trabajo es útil como mapa del terreno. Muestra que el bucle no es una sola técnica, sino un conjunto de decisiones con distintos compromisos, y sugiere cuál elegir según su presupuesto. Por ahora es material para quienes entrenan modelos por sí mismos, pero la tendencia apunta claramente a exprimir más de cada parámetro.