Pretraining Recurrent Networks without Recurrence

Resumen Técnico: Preentrenamiento de Redes Recurrentes sin Recurrencia

Planteamiento del Problema

El entrenamiento de Redes Neuronales Recurrentes (RNN) requiere fundamentalmente resolver el problema de la asignación de crédito a través de secuencias largas: determinar cómo los inputs tempranos influyen en los outputs futuros cuando intervienen múltiples computaciones intermedias. La solución estándar, la Propagación de Gradiente a Través del Tiempo (BPTT - Backpropagation Through Time), despliega la RNN en el tiempo y propaga los gradientes hacia atrás a través de todo el grafo de computación. Este enfoque sufre dos limitaciones críticas:

1. Cuello de Botella Secuencial: BPTP es inherentemente secuencial en el tiempo, lo que impide la paralelización y limita la escalabilidad en el hardware moderno.
2. Inestabilidad del Gradiente: La longitud del camino del gradiente es $O(T)$ (donde $T$ es la longitud de la secuencia). Esto provoca gradientes que se desvanecen o explotan, dificultando el aprendizaje de dependencias de largo alcance y causando un "sesgo de recencia" donde el modelo olvida la información temprana.

Si bien los Transformers resolvieron el problema del paralelismo, carecen de una memoria comprimida y de tamaño fijo del pasado, almacenando en su lugar todo el historial de tokens. Las RNN lineales ofrecen memoria fija y paralelismo, pero están limitadas por funciones de transición lineales, lo que restringe su expresividad para tareas secuenciales complejas como el seguimiento de estados.

Metodología: Entrenamiento de Memoria Supervisada (SMT)

El artículo propone el Entrenamiento de Memoria Supervisada (SMT - Supervised Memory Training), un método para entrenar RNN no lineales que evita por completo la propagación de la recurrencia del crédito. SMT reduce el entrenamiento de la RNN al aprendizaje supervisado sobre etiquetas de transición de memoria de un solo paso.

Mecanismo Central

SMT desacopla qué recordar (representación de la memoria) de cómo actualizar la memoria (dinámica de la memoria):

1. Modelo Maestro (Aprendizaje de Estado Predictivo): Se entrena un par codificador-decodificador basado en Transformer, que es paralelo en el tiempo, para comprimir el contexto pasado en un estado de memoria $m_t$ que es un estado predictivo. Este estado retiene solo la información necesaria para predecir salidas futuras, descartando detalles irrelevantes. El codificador mapea el contexto $x_{\le t}$ al estado $m_t$, y el decodificador predice salidas futuras $y_{\ge t}$ dado $m_t$ y entradas futuras.
2. Modelo Estudiante (Dinámica de la RNN): La RNN $f_\theta$ es entrenada para imitar las transiciones de memoria del maestro. En lugar de desplegar la RNN, aprende la función de actualización de un solo paso: $(m_t, x_{t+1}) \to \hat{m}_{t+1}$.
3. Funciones de Pérdida:
   • Pérdida de Decodificación ($L_{dec}$): Entrena al maestro para predecir el futuro a partir de la memoria comprimida.
   • Pérdida de Dinámica ($L_{dyn}$): Entrena a la RNN para predecir el siguiente estado de memoria $\hat{m}_{t+1}$ dado el estado de memoria actual $m_t$ y la entrada $x_{t+1}$, supervisado por la verdad de campo (ground truth) de la memoria $m_{t+1}$ del maestro.
   • Pérdida de Uniformidad ($L_{unif}$): Evita que el espacio de memoria colapse en un único punto.

Pipeline de Entrenamiento

1. Preentrenamiento (SMT): El codificador-decodificador del maestro y la RNN se entrenan conjuntamente (o secuencialmente) utilizando las pérdidas anteriores. Crucialmente, la RNN nunca se despliega durante esta fase. El camino del gradiente entre dos tokens cualesquiera es $O(1)$ porque el crédito se asigna mediante el objetivo de estado predictivo del maestro, no mediante la retropropagación a través de los pasos recurrentes de la RNN.
2. Ajuste Fino (DMT - Entrenamiento de Memoria DAgger): Para abordar el desajuste entre el entrenamiento y la prueba (train-test mismatch) donde la RNN utiliza sus propias memorias predichas (que pueden desviarse de la memoria de verdad de campo del maestro) en lugar de las memorías reales del maestro, se introduce una fase ligera de DAgger. La RNN se despliega con sus propios estados predichos, y la pérdida de dinámica se calcula para corregir esta deriva mediante aprendizaje de imitación on-policy.

Contribuciones Clave

• Entrenamiento en Paralelo Temporal de RNN No Lineales: SMT permite el entrenamiento de RNN no lineales expresivas sin despliegue secuencial, logrando una longitud de camino de gradiente de $O(1)$ entre tokens.
• Asignación de Crédito Estable a Largo Alcance: Al reducir el problema al aprendizaje supervisado sobre estados predictivos, SMT elimina los problemas de desvanecimiento o explosión de gradientes asociados con caminos de $O(T)$, permitiendo un aprendizaje estable de dependencias de largo alcance.
• Inferencia de Memoria Fija: A diferencia de los Transformers, las RNN entrenadas con SMT mantienen un tamaño de memoria fijo ($O(M)$) durante la inferencia, lo que las hace adecuadas para secuencias ilimitadas (por ejemplo, aprendizaje de por vida).
• Desacoplamiento de Representación y Dinámica: El método separa el aprendizaje de la compresión de memoria (gestionado por el maestro paralelo) de la dinámica de la memoria (gestionada por la RNN), simplificando el panorama de optimización para la RNN.

Resultados Experimentales

Los autores evaluaron SMT frente a baselines de BPTT y Transformers en tareas sintéticas y conjuntos de datos del mundo real:

• Tareas Sintéticas: En tareas que prueban la estabilidad del gradiente (recuperación), capacidad de memoria (copia de cadenas), seguimiento de estados (operaciones de pila), recuperación asociativa y aprendizaje en contexto, SMT→DMT superó consistentemente a BPTT. BPTT tuvo dificultades significativas a medida que aumentaba la longitud de la secuencia, mientras que el rendimiento de SMT se mantuvo estable.
• Modelado de Secuencias de Píxeles (Tarea de Attneave): En los conjuntos de datos MNIST y Sketchy (modelados como secuencias de píxeles 1D), las RNN entrenadas con BPTT fallaron en capturar dependencias estructurales de largo alcance (por ejemplo, la continuidad del trazo). Las RNN de SMT→DMT generaron imágenes coherentes con éxito, demostrando una capacidad superior para integrar información a lo largo de cientos de pasos de tiempo.
• Eficiencia: SMT y SMT→DMT requirieron significativamente menos FLOPs secuenciales (una medida de computación serial) que BPTT para lograr un rendimiento comparable o superior, particularmente en tareas que requieren memoria de largo alcance. La eficiencia de datos fue comparable a la de BPTT en el modelado de lenguaje, pero superior en el modelado de píxeles.
• Leyes de Escala: SMT→DMT exhibió mejoras de rendimiento suaves con el aumento de la longitud del contexto, el tamaño de la memoria y los parámetros del modelo. Notablemente, el rendimiento de la RNN se acercó al rendimiento del Transformer maestro a escalas mayores.
• Generalización: Las RNN entrenadas con SMT generalizaron mejor que sus maestros Transformer a longitudes de secuencia más largas que las observadas durante el entrenamiento, lo que sugiere un sesgo inductivo más robusto para la abstracción temporal.

Significado y Reivindicaciones

El artículo afirma que SMT ofrece un camino para desbloquear la escala de modelos que construyen abstracciones temporales de la experiencia pasada. Al permitir que las RNN no lineales aprendan dependencias de largo alcance con gradientes estables y memoria fija, SMT aborda las limitaciones tanto de BPTT (inestabilidad, secuencialidad) como de los Transformers (crecimiento ilimitado de la memoria).

Los autores posicionan a SMT principalmente como un algoritmo de preentrenamiento. Reconocen que, si bien SMT aprende eficazmente las representaciones de memoria, es necesaria una fase de post-entrenamiento ligera (como DMT) para mitigar la deriva y adaptar el modelo a tareas específicas de downstream, permitiendo potencialmente que el modelo supere los límites de expresividad del codificador maestro. El trabajo sugiere que combinar el entrenamiento en paralelo temporal con la inferencia de memoria fija y la expresividad máxima es una dirección viable para futuros modelos de secuencias, permitiendo potencialmente que los agentes aprendan de horizontes ilimitados como la experiencia de una vida humana.