Resumen Técnico: Leviathan – Desacoplamiento de Representaciones de Entrada y Salida en Modelos de Lenguaje

1. Enunciado del Problema

Los modelos de lenguaje grandes (LLM) modernos suelen emplear una arquitectura de pesos atados, donde una única matriz de incrustaciones $E \in \mathbb{R}^{V \times D}$ cumple una doble función: mapear tokens de entrada a representaciones continuas y proyectar estados ocultos de vuelta al espacio del vocabulario para la predicción de salida. Este acoplamiento obliga a una única matriz a satisfacer simultáneamente dos objetivos distintos:

1. Representación de Entrada: Requiere suavidad, donde tokens semánticamente relacionados ocupan regiones cercanas en el espacio de representación.
2. Clasificación de Salida: Requiere discriminación, separando efectivamente todos los tokens a través del vocabulario.

Los autores hipotetizan que este acoplamiento impone una limitación arquitectónica. Si bien los modelos completamente desatados (que aprenden matrices de entrada y salida separadas) desacoplan teóricamente estos objetivos, demuestran ser ineficaces en la práctica. Los experimentos muestran que los modelos completamente desatados tienen un rendimiento inferior al de las líneas base atadas a pesar de un aumento de ~50% en la cantidad de parámetros, lo que sugiere que los parámetros adicionales no se utilizan de manera efectiva debido a dificultades de optimización y una señal de entrenamiento insuficiente para incrustaciones discretas, particularmente para tokens raros.

2. Metodología: La Arquitectura Leviathan

Leviathan introduce una arquitectura Transformer que desacopla la representación de entrada de la predicción de salida sin incurrir en el costo prohibitivo de parámetros de las incrustaciones completamente desatadas.

2.1 Vectorización de Incrustaciones Aprendidas (LEV)

La innovación central es la capa LEV, que reemplaza la matriz de incrustación de entrada densa con una función continua aprendida $G: \{0, \dots, V-1\} \to \mathbb{R}^D$. Esta función mapea índices de tokens a incrustaciones a través de un mecanismo compacto de compartición de parámetros:

• Indexación Composicional: El vocabulario se factoriza en $k$ componentes utilizando una descomposición en base-$b$, convirtiendo un índice de token $i$ en una tupla $(i_1, \dots, i_k)$.
• Libros de Códigos Compartidos: Estos índices se mapean a libros de códigos compartidos $C_1, \dots, C_k$, generando una representación semilla $z(i)$. Esto permite que los tokens raros se beneficien de las señales de gradiente pasadas a través de entradas compartidas utilizadas por tokens estructuralmente similares.
• Proyección Continua: La semilla se proyecta y normaliza a una coordenada latente acotada $\tilde{z}_\ell \in [0, 1]^{d_{seed}}$.
• Expansión de Splines: Cada dimensión se modela mediante una expansión de base de B-splines cuadráticos univariados con $\kappa nodos.
• Producto Tensorial: Las dimensiones se combinan mediante un producto tensorial de rango-$r$.
• Suma de Cabezas: La incrustación final es la suma de las salidas de $h$ cabezas paralelas después de una proyección lineal.

2.2 Arquitectura y Desacoplamiento

• Cabeza de Salida Desatada: La capa LEV permite que la cabeza de proyección de salida permanezca completamente desatada (una matriz densa independiente de la entrada).
• Eficiencia de Parámetros: Dado que la capa LEV es significativamente más pequeña que una tabla de incrustaciones densa, la sobrecarga de la cabeza de salida desatada es negligible. A la escala de 1.2 mil millones de parámetros, el aumento total de parámetros es solo del 0.2%.
• Sesgo Inductivo: La arquitectura impone un sesgo inductivo de suavidad sobre el espacio de tokens, en contraste con la geometría discreta y potencialmente anisotrópica de las tablas de búsqueda estándar.

3. Configuración Experimental

• Modelos: Los experimentos se realizaron en Transformers solo-decodificador en tres escalas: 200M, 400M y 1.2B parámetros.
• Líneas Base:
  • ATADO (TIED): Línea base estándar de pesos atados.
  • DESATADO (UNTIED): Modelo completamente desatado (611M parámetros a escala de 400M).
  • ALBERT: Factorización coincidente de parámetros de entrada (compresión lineal) con una cabeza de salida desatada.
  • Ablaciones: MLP plano y MLP factorizado coincidentes en parámetros para aislar las contribuciones de la no linealidad y la estructura composicional.
• Entrenamiento: Todos los modelos fueron entrenados en The Pile (sin derechos de autor) con hiperparámetros idénticos, semillas y presupuestos de tokens (hasta 24.1 mil millones de tokens para modelos de 1.2B).
• Evaluación: Rendimiento de modelado de lenguaje en un conjunto de validación retenido de The Pile, generalización zero-shot en Wikitext-103, y benchmarks aguas abajo (HellaSwag, ARC, PIQA, WinoGrande, LAMBADA).

4. Resultados Clave

4.1 Rendimiento de Modelado de Lenguaje

Leviathan supera consistentemente a la línea base de pesos atados en todas las escalas, ampliándose la brecha de rendimiento a medida que avanza el entrenamiento y aumenta la escala del modelo.

• Escala 1.2B: Leviathan redujo la perplejidad de validación en un 9.2% (0.087 nats) en comparación con la línea base atada.
• Eficiencia de Muestra: A escala de 1.2B, Leviathan requirió 2.1 veces menos tokens de entrenamiento para alcanzar la pérdida final de la línea base atada.
• Dinámicas de Entrenamiento: En escalas más pequeñas (200M), Leviathan inicialmente tuvo un rendimiento inferior antes de superar la línea base; a 1.2B, superó a la línea base durante todo el entrenamiento.

4.2 Comparación con Alternativas

• Vs. Completamente Desatado: El modelo completamente desatado (UNTIED) tuvo un rendimiento inferior a la línea base atada durante todo el entrenamiento, confirmando que el desacoplamiento ingenuo sin parametrización eficiente es subóptimo.
• Vs. Factorización Lineal: La línea base factorizada al estilo ALBERT también tuvo un rendimiento inferior a la línea base atada, lo que indica que la compresión lineal por sí sola no puede recuperar los beneficios representacionales de la geometría atada.
• Vs. Leviathan: Leviathan superó tanto a UNTIED como a ALBERT, demostrando que la generación continua no lineal es superior tanto al desatado ingenuo como a la factorización lineal dentro de presupuestos de computación óptimos según Chinchilla.

4.3 Transferencia Aguas Abajo y Generalización

• Tareas Aguas Abajo: A escala de 1.2B, Leviathan mejoró el rendimiento en las seis evaluaciones de benchmarks aguas abajo. Destacablemente, logró una reducción del 30% en la perplejidad en LAMBADA.
• Transferencia Zero-Shot: En Wikitext-103 (una distribución diferente), Leviathan mostró mejoras de perplejidad crecientes con la escala (2.4% a 200M, 12.3% a 1.2B), lo que sugiere que sus representaciones continuas se transfieren de manera más efectiva a través de distribuciones de datos.

4.4 Análisis Estratificado por Frecuencia

Las ganancias de rendimiento están fuertemente concentradas en la cola larga de la distribución de tokens:

• Tokens Raros: Para tokens que aparecen menos de 1 en 10 millones de veces, Leviathan redujo la perplejidad en un 81% a escala de 1.2B.
• Tokens Frecuentes: Para tokens que aparecen más de 1 en 10 mil veces, la mejora cayó a casi cero (y ligeramente negativa para los más frecuentes), ya que las incrustaciones discretas son suficientes cuando las señales de entrenamiento son fuertes.
• Mecanismo: La parametrización continua comparte gradientes a través de tokens estructuralmente similares, proporcionando una señal robusta donde las búsquedas discretas fallan debido a la escasez de datos.

4.5 Robustez

• Tamaño de Vocabulario: Los resultados se mantuvieron consistentes en tamaños de vocabulario de 50k, 100k y 200k, con mejoras relativas de ~2% en general.
• Ablaciones: Las ganancias se descompusieron en: representación continua (+1.13%), estructura composicional (+0.37%) y suavidad de spline separable (+0.62%).

5. Significado y Afirmaciones

El artículo afirma que la práctica estándar de atar pesos impone una limitación arquitectónica que no se resuelve simplemente desatando los pesos. En cambio, la solución óptima implica desacoplar las representaciones de entrada y salida mientras se impone suavidad mediante parametrización continua.

• Insight Arquitectónico: La geometría representacional de las incrustaciones de entrada importa. La parametrización continua es un medio eficiente en datos y computación para mejorar esta geometría, particularmente para tokens raros donde las incrustaciones discretas luchan.
• Eficiencia: Leviathan logra estas ganancias con una sobrecarga de parámetros negligible (0.2% a escala de 1.2B), lo que la convierte en una actualización práctica para los modelos de lenguaje modernos.
• Escalabilidad: Los beneficios de la capa LEV se compilan con la escala y los datos, lo que sugiere que a medida que los modelos crecen, la ventaja de las representaciones de entrada suaves y continuas sobre las búsquedas discretas se vuelve más pronunciada.

Los autores concluyen que, aunque el patrón de mejora dependiente de la frecuencia es correlacional, la consistencia a través de escalas, vocabularios y variantes arquitectónicas proporciona evidencia robusta de que la parametrización continua mitiga eficazmente las dificultades de optimización asociadas con el aprendizaje de incrustaciones discretas para tokens raros.