Investigadores encuentran en Pythia una ventana del 0,4–1,1% del preentrenamiento en la que se reorganiza el readout

Se suele dar por hecho que un modelo ha adquirido una habilidad en el momento en que empieza a producir el token correcto. Un nuevo trabajo, publicado en Hugging Face Daily Papers, muestra que esa lógica introduce un desfase. Los estados ocultos pueden contener información sobre la tarea antes de que la salida del modelo la exprese. Con Pythia como ejemplo, los autores identificaron una ventana estrecha, del 0,4–1,1% del recorrido del preentrenamiento, en la que se reorganiza la matriz de unembedding (readout) y en cuyas filas aparecen las categorías de WordNet.

La salida también aprende

La predicción de un token pasa por la matriz de unembedding, que convierte el estado oculto en logits sobre el vocabulario. Esta matriz se entrena a la vez que el resto del modelo. Por eso, una gráfica de comportamiento en la que la precisión sube a partir del paso N mezcla dos procesos: lo que el modelo ya ha codificado internamente y lo que su readout ha aprendido a expresar hacia fuera.

Los autores lo comprobaron de forma directa. Las sondas lineales sobre los estados ocultos de checkpoints tempranos ya decodifican información de la tarea, aunque las propias salidas del modelo todavía no la muestran. Existe, por tanto, una brecha entre el conocimiento y su manifestación, y es esa brecha la que distorsiona la datación de muchas «transiciones» en el desarrollo del modelo.

Parameter-trajectory crosscoding

Para seguir el readout, los investigadores propusieron parameter-trajectory crosscoding. El método entrena un único diccionario disperso sobre las filas del readout de cada token en cada checkpoint. Los crosscoders clásicos trabajan con activaciones y exigen pasar datos por el modelo. Aquí solo hacen falta los pesos de los checkpoints, lo que simplifica notablemente el análisis de modelos abiertos que conservan su historial de entrenamiento.

Principales observaciones:

• ventana de reorganización del readout en Pythia: 0,4–1,1% del preentrenamiento;
• en esa ventana, las categorías de WordNet pasan a ser decodificables a partir de las filas del readout;
• hacia el final de la ventana, las sondas sobre los estados ocultos ya leen información que no está en las salidas;
• el análisis solo requiere los pesos de los checkpoints, sin activaciones;
• con un learning rate más alto en el readout, el loss final apenas cambia.

Un experimento de control con el learning rate

La correlación no basta, así que los autores introdujeron una intervención. Entrenaron desde cero modelos pequeños que se diferenciaban únicamente en el learning rate del readout. Si la brecha depende de la velocidad de aprendizaje de la capa de salida, debería moverse junto con ella.

Así ocurrió. Con un learning rate más alto, el readout se forma antes y la brecha entre el estado oculto y la salida se cierra más deprisa. El loss final apenas se resiente. Es decir, la velocidad de aprendizaje del readout influye en cuándo se hace visible una capacidad desde fuera, pero no en el punto al que llega el modelo al final.

Cómo aplicar el resultado en la práctica

La conclusión se refiere a la fase temprana del entrenamiento, donde se sitúan muchas transiciones «repentinas». Si construyes curvas de aparición de capacidades, una sola métrica sobre las salidas no es suficiente. Conviene añadir sondas lineales sobre los estados ocultos y comparar las fechas: si la sonda se adelanta a la salida, lo más probable es que estés midiendo la maduración del readout y no la aparición de la habilidad.

Para quienes entrenan sus propios modelos, también es una razón para plantearse un learning rate específico para el unembedding. El trabajo se ha validado en modelos pequeños y en Pythia, así que hay que ser cautos al extrapolar las conclusiones a arquitecturas modernas de gran tamaño. Aun así, la herramienta ya es aplicable a checkpoints abiertos, porque los pesos están disponibles en todos ellos.