Duplicar capas durante el entrenamiento elevó GSM8K del 12% al 19% en un modelo de 3B

Los investigadores Ferdinand Kapl, Kaitlin Maile, Stefan Bauer y Johannes von Oswald compararon tres formas de construir un modelo de lenguaje de hasta 3.000 millones de parámetros. La primera es una red densa convencional.

La segunda es una red que se va haciendo más profunda durante el entrenamiento mediante la copia de las capas intermedias (el método LIDAS). La tercera es una red «en bucle», en la que las mismas capas se ejecutan varias veces en una sola pasada. El estudio se centró en el aprendizaje en contexto (in-context learning, ICL): la capacidad del modelo para entender una tarea a partir de un par de ejemplos en el prompt, sin reentrenamiento.

Resultó que tanto el crecimiento en profundidad como el bucle aprenden ICL notablemente mejor que el modelo denso. A escala de 3B, LIDAS alcanzó cerca de un 19% de coincidencia exacta en GSM8K con 8 ejemplos en el prompt; el modelo denso, cerca de un 12%.

Dos formas de reutilizar la profundidad

En un modelo convencional, el número de capas se fija de antemano y cada capa se entrena desde cero. En LIDAS ocurre lo contrario: el entrenamiento arranca con una red poco profunda, después se duplican las capas intermedias y la red crece sobre la marcha. Al principio las copias son casi idénticas; luego divergen. La fase inicial del entrenamiento se hace con un modelo pequeño, por lo que resulta más barata. Según los autores, en 3B LIDAS gasta un 26% menos de FLOPs (operaciones de coma flotante) en el entrenamiento y, aun así, supera al modelo denso.

El bucle es más sencillo. Se toma un conjunto reducido de capas y los datos pasan por él varias veces en una sola pasada hacia delante. El resultado son menos parámetros únicos y una profundidad efectiva mayor.

Los experimentos se realizaron con modelos de entre 360M y 3B de parámetros. El código se basa en la biblioteca nanotron de Hugging Face. Los autores no publican código propio ni pesos.

Dónde ganan los modelos

GSM8K es un conjunto de problemas de matemáticas de primaria en formato de texto. Se incluyen ocho ejemplos resueltos en el prompt y se comprueba si la respuesta del modelo coincide con la de referencia. Resultados de los modelos de 3B:

• LIDAS: cerca de un 19% de coincidencias exactas;
• modelo denso: cerca de un 12%.

La segunda prueba es RegBench: lenguajes formales nuevos que el modelo no ha visto antes. A partir de unos pocos ejemplos, debe captar las reglas y predecir el siguiente símbolo. La calidad se mide con la distancia de variación total, es decir, la diferencia entre las probabilidades predichas y las reales. Cuanto menor, mejor:

• crecimiento en profundidad: ~0.119;
• bucle: ~0.125;
• modelo denso: ~0.137.

Aquí se mide ICL en estado puro: en la prueba no hay nada que recordar de los pesos, solo queda leer el contexto.

Cabezas de inducción multiplicadas y el precio de la mejora

Los autores explican el resultado mediante las cabezas de inducción. Son cabezas de atención que localizan en el contexto una aparición anterior del token actual y sugieren lo que venía después. Se consideran uno de los mecanismos principales del ICL.

En el modelo de 3B con LIDAS, 10 de las 12 cabezas de inducción más fuertes proceden de dos ancestros de la red inicial de cuatro capas. El circuito encontrado en el modelo pequeño se copia a lo largo de toda la profundidad a medida que la red crece. Parece que luego no lo desplaza el conocimiento que la red memoriza en los pesos.

Otra comprobación es la desactivación. Los autores apagaron el 20% superior de estas cabezas y midieron cuánto cae el F1 medio en 19 tareas con varios ejemplos:

• modelo denso: caída de aproximadamente 0.08;
• LIDAS: aproximadamente 0.18;
• modelo en bucle: aproximadamente 0.28.

Es decir, los modelos con profundidad ampliada y los modelos en bucle dependen más de estos circuitos, y en ellos se concentra toda su capacidad de aprender a partir de ejemplos.

Qué deben tener en cuenta los ingenieros

Para quienes preentrenan modelos, la conclusión es práctica. Si se busca ahorrar cómputo y obtener un buen ICL, el crecimiento gradual de la profundidad parece una opción razonable. En 3B ofreció mejores resultados y un 26% de ahorro en FLOPs.

El bucle tiene un coste mayor. Este modelo tiene menos parámetros únicos y resulta peor a la hora de extraer los conocimientos almacenados en los pesos. La mejora en ICL se paga: el modelo lee mejor el contexto, pero recuerda peor los hechos.

También hay una salvedad sobre la fiabilidad. Cada configuración, de 360M a 3B, se entrenó con una sola semilla, es decir, con un único estado inicial del generador de números aleatorios. La dispersión entre ejecuciones no se ha medido. Siete puntos porcentuales en GSM8K parecen convincentes, pero sin repeticiones no se puede decir cuán estables son. Además, un tamaño de hasta 3B es pequeño para los estándares de los modelos actuales, y trasladar las conclusiones a decenas de miles de millones de parámetros requiere una verificación aparte.

Aun así, el resultado es inesperado. El modelo denso, entrenado de la forma habitual, salió peor justo donde se lo consideraba fuerte: en la capacidad de aprender de los ejemplos del contexto. La explicación de los autores, la copia temprana de circuitos de atención, se puede poner a prueba con otros tamaños.