Resumen Técnico: Predicción de Latentes Jerárquicos para Modelos de Lenguaje

Planteamiento del Problema

La predicción estándar de siguiente token (NTP, por sus siglas en inglés) con entrenamiento de forzado de profesor ha permitido el desarrollo de modelos de lenguaje transformer altamente capaces. Sin embargo, este paradigma introduce un desajuste entre el entrenamiento y la inferencia conocido como sesgo de exposición (exposure bias). Durante la inferencia, la generación autorregresiva depende de las salidas previas del propio modelo en lugar de la verdad de referencia, lo que genera errores acumulativos que degradan la calidad de la generación a largo plazo.

Intentos recientes para mitigar esto, como la Predicción de Múltiples Tokens (MTP) y la Predicción de Siguiente Latente (NextLat), introducen objetivos auxiliares para predecir futuros tokens o estados latentes. No obstante, estos enfoques enfrentan limitaciones:

• MTP suele sufrir de errores acumulativos durante los despliegues (rollouts) de múltiples pasos.
• NextLat impone la autopredicción en un espacio latente pero, típicamente, predice solo un paso hacia adelante. Aunque es efectivo para la consistencia de transición local, proporciona una presión de gradiente solo indirecta para capturar la estructura en escalas temporales más largas, dejando las dependencias de largo horizonte atenuadas por el despliegue de la dinámica secuencial.

El desafío central es cómo moldear la representación de un modelo de lenguaje para codificar estructuras predictivas multiescala durante el preentrenamiento sin incurrir en sobrecarga arquitectónica o latencia de inferencia.

Metodología: Predicción de Latentes Jerárquicos (HiLP)

Los autores proponen la Predicción de Latentes Jerárquicos (HiLP), un método de entrenamiento que introduce jerarquía temporal en el objetivo de autopredicción en el espacio latente. La arquitectura consta de cuatro mapas parametrizados aprendidos conjuntamente: un cabezal de salida ($p_\theta$), dinámica de nivel inferior ($p_\psi$), dinámica de nivel superior ($p_\phi$) y un cabezal de NTP combinado ($p_\rho$).

Componentes de la Arquitectura

1. Latentes de Nivel Inferior ($h_t$): Un transformer de base $G_\theta$ produce latentes de nivel inferior $h_t$ a partir del prefijo de tokens $X_{1:t}$.
2. Latentes Abstractos de Nivel Superior ($u_t$): Un módulo de atención de ventana deslizante (SWA) con un tamaño de ventana $k$ agrega los últimos $k$ estados de nivel inferior en un latente abstracto de nivel superior: $u_t = \text{SWA}(h_{t-k+1:t})$.
3. Dinámica de Nivel Superior: Un modelo $p_\phi$ predice $k$ pasos hacia adelante en el espacio abstracto ($\hat{u}_{t+k}$) sin entrada de tokens, fomentando la planificación con visión de futuro más allá de las transiciones de un solo paso.
4. Cabezal de NTP Combinado: La predicción del siguiente token se condiciona tanto al latente de nivel inferior $h_t$ como al latente de nivel superior $\tilde{u}_t$ (donde $\tilde{u}_t$ incluye un corte de gradiente para evitar que el cabezal combinado actualice los latentes inferiores a través de la ventana SWA).

Objetivos de Entrenamiento

El loss total es una suma ponderada de cinco términos:

1. Loss de NTP Estándar ($L_{ntp}$): Entropía cruzada estándar para la predicción del siguiente token.
2. Consistencia de Transición de Nivel Inferior ($L_h$): Pérdida SmoothL1 que asegura que el modelo de dinámica de nivel inferior $p_\psi$ prediga con precisión el siguiente estado latente de nivel inferior.
3. Divergencia KL ($L_{KL}$): Asegura que la distribución de los latentes predichos coinciera con la distribución objetivo.
4. Consistencia de Transición de Nivel Superior ($L_u$): Pérdida SmoothL1 que asegura que el modelo de dinámica de nivel superior $p_\phi$ prediga con precisión el latente abstracto $k$ pasos adelante.
5. Loss de NTP Combinado ($L_{cntp}$): Entropía cruzada para el cabezal combinado, aprovechando tanto $h_t$ como $\tilde{u}_t$.

Estrategia de Inferencia: Una característica de diseño crítica de HiLP es que todo el aparato jerárquico (módulo SWA, modelos de dinámica y cabezal combinado) se descarta en el momento de la inferencia. Solo se utiliza el cabezal de NTP estándar, lo que garantiza cero sobrecarga arquitectónica y ningún aumento en la latencia de inferencia.

Contribuciones Clave

1. Entrenamiento de Representación Jerárquica: La introducción de HiLP, que integra el aprendizaje de autopredicción multiescala en el preentrenamiento de transformers mediante atención de ventana deslizante sobre latentes, un modelo de dinámica de nivel superior y un cabezal de NTP combinado.
2. Evaluación: Benchmarking exhaustivo de los modelos resultantes en precisión de tareas descendentes (codificación y razonamiento de múltiples pasos), comparando HiLP contra MTP y NextLat.
3. Intervención Pura en Tiempo de Entrenamiento: Demostración de que la compleja estructura jerárquica sirve puramente como una señal auxiliar de entrenamiento, permitiendo que el modelo mantenga la eficiencia de la decodificación autorregresiva estándar durante el despliegue.

Resultados Experimentales

Los autores entrenaron modelos de 1B de parámetros en 100B de tokens utilizando 8× GPUs NVIDIA B200.

Rendimiento en Benchmarks

• Codificación (HumanEval): HiLP alcanzó una puntuación pass@1 de 11.33, superando a NTP estándar (8.77), MTP (9.21) y NextLat (10.58).
• Razonamiento de Múltiples Pasos: En benchmarks simbólicos y de múltiples pasos (DataComp para LLMs), HiLP demostró un rendimiento mejorado, atribuido a la predicción de latentes multiescala.

Decodificación Especulativa

HiLP mostró una eficiencia superior en tareas de decodificación especulativa:

• Datos de Código: HiLP logró un promedio de 3.41 tokens aceptados por paso de redacción, comparado con 2.57 para MTP y 3.26 para NextLat.
• Datos de Nemotron-ClimbMix: HiLP logró 2.59 tokens aceptados, superando ligeramente a NextLat (2.57) y MTP (2.11).
• Estabilidad del Despliegue de Latentes: El análisis de las curvas de entropía cruzada de los latentes reveló que, si bien HiLP mantiene una pérdida de predicción a corto plazo comparable a la de NextLat, reduce significativamente el error en horizontes más largos (hasta 8 pasos), lo que indica despliegues de latentes de múltiples pasos más estables.

Eficiencia

Aunque HiLP añade parámetros auxiliares (un total de 1.27B frente a 1.06B para el backbone), la ruta de inferencia permanece idéntica a la del modelo base. El rendimiento de entrenamiento (throughput) para HiLP fue aproximadamente del 65% respecto al baseline de NTP, lo cual es mayor que el de MTP (49%) pero menor que el de NextLat (~83%), reflejando el costo del grafo de pérdida jerárquico adicional.

Significado y Reivindicaciones

El artículo afirma que HiLP aborda con éxito la "brecha miópica" en la predicción del siguiente token al introducir la jerarquía temporal directamente en el objetivo de autopredicción en el espacio latente. A diferencia de modelos jerárquicos previos que consumen complejidad en el tiempo de inferencia (por ejemplo, Large Concept Models o Byte Latent Transformers), HiLP moldea la representación durante el preentrenamiento para codificar una estructura predictiva multiescala, la cual se elimina por completo para la inferencia.

Los autores postulan que este enfoque permite al modelo aprender dinámicas de transición de latentes más ricas y capturar mejor las dependencias de largo horizonte sin sacrificar la velocidad o la simplicidad de la decodificación autorregresiva estándar a nivel de token. Este trabajo se posiciona en la intersección de los objetivos de preentrenamiento más allá de NTP, el modelado de lenguaje en el espacio latente y los modelos jerárquicos, ofreciendo un método novedoso para mejorar las capacidades de razonamiento y planificación de largo alcance en los modelos de lenguaje.