OLLM: Options-based Large Language Models

1. Problema Identificado

Los Modelos de Lenguaje Grandes (LLM) tradicionales modelan la distribución de un único "siguiente token" y dependen de heurísticas de muestreo (como temperatura o top-k/p) para generar diversidad. Sin embargo, este enfoque presenta limitaciones:

• Ineficiencia en la exploración: Buscar la mejor salida en el espacio completo del vocabulario es costoso y poco eficiente en términos de muestras.
• Falta de control explícito: Las variaciones en la generación a menudo surgen de manera aleatoria en lugar de ser controladas por una política deliberada.
• Problemas de alineación: Los métodos de ajuste fino por refuerzo (RLHF) que operan sobre todo el vocabulario pueden provocar desalineaciones, como cambios de idioma o razonamiento degenerado, y a menudo requieren penalizaciones complejas (como KL-divergencia) o etiquetas adicionales.
• Multimodalidad no capturada: En tareas como el razonamiento matemático, existen múltiples continuaciones plausibles en un mismo paso, pero el modelo estándar compite por masa de probabilidad dentro de un solo softmax, lo que puede llevar a la sobreajuste a un modo dominante.

2. Metodología: OLLM (Options-based Large Language Models)

Los autores proponen OLLM, un método general que reemplaza la predicción de un solo token por un conjunto de "opciones" aprendidas para el siguiente token, indexadas por una variable latente discreta.

Arquitectura

OLLM actúa como un "plug-in" ligero que se inserta antes de la capa de salida (lm_head) de cualquier LLM preentrenado, sin modificar el backbone original. Consta de dos módulos adicionales:

1. Codificador (Encoder): Durante el entrenamiento, toma la representación del token objetivo y el estado oculto contextual para producir una distribución sobre un espacio latente discreto pequeño ($Z$). Esto evita escalar con el tamaño del vocabulario.
2. Decodificador (Decoder): Toma el estado oculto y una muestra del espacio latente ($z$) para generar un vector de sesgo ($\Delta h_t$) que se suma al estado oculto antes de la proyección al vocabulario.

Entrenamiento y Aprendizaje de Políticas

• Fase de SFT (Supervised Fine-Tuning): Se entrena el codificador y el decodificador para predecir el token correcto dado el contexto y la opción latente. Se utiliza una pérdida de entropía cruzada regularizada con una divergencia KL hacia una distribución uniforme para evitar el colapso de las latentes.
• Política en Espacio Latente: Se entrena una política ligera ($\pi_\phi$) que, durante la inferencia, selecciona la opción latente adecuada basándose únicamente en el estado oculto del LLM.
• Desventaja de Opciones: En lugar de que todos los tokens compitan en un solo espacio de probabilidad, OLLM descompone la distribución en una mezcla sobre las opciones latentes. Esto permite que continuaciones alternativas se optimicen bajo diferentes latentes, desacoplando la competencia.

3. Contribuciones Clave

• Modelado Explícito de Variación: OLLM hace explícita la variación en un manifold latente compacto, permitiendo que una política de bajo dimensión controle la generación en lugar de operar sobre el vocabulario completo.
• Eficiencia de Parámetros: Es un método extremadamente eficiente; en los experimentos, solo el 1.56% de los parámetros del modelo (1.7B) son entrenables, manteniendo el backbone congelado.
• Sin Etiquetas Externas: A diferencia de trabajos previos que requieren etiquetas de atributos (sentimiento, tema), OLLM aprende las opciones de manera no supervisada a partir de la estructura del texto y el contexto.
• Alineación Estructural: La alineación y la reducción de comportamientos indeseados (como cambios de idioma) surgen de la estructura del modelo (restringido a opciones aprendidas durante SFT) en lugar de depender de pérdidas de alineación manuales o penalizaciones KL adicionales.

4. Resultados

Los autores evaluaron OLLM utilizando un modelo base de 1.7B parámetros (distilado de Qwen2.5/Deepseek R1) entrenado en el conjunto de datos OpenMathReasoning y evaluado en OmniMath.

• Rendimiento de Línea Base: Los modelos adaptados con LoRA (State-of-the-Art) alcanzaron un máximo de 51% de precisión en la respuesta final.
• Rendimiento de OLLM:
  • Bajo la selección óptima de la variable latente (el "oráculo"), OLLM alcanzó un ~70% de precisión en la respuesta final.
  • La precisión por token mejoró hasta un ~83%.
• Eficiencia de Muestreo: Las políticas simples operando en el espacio latente de baja dimensión lograron mejoras sustanciales en la precisión final, demostrando que separar las opciones plausibles evita el sobreajuste a un modo dominante.
• Análisis Cualitativo: Se observó que OLLM aprende opciones genuinamente diversas cuando hay ambigüedad (ej. diferentes direcciones de razonamiento lógico) y colapsa a una elección consistente en casos deterministas (ej. secuencias de símbolos matemáticos fijos).

5. Significado e Impacto

El trabajo demuestra que el modelado de "siguiente token" basado en opciones mejora significativamente la controlabilidad, robustez y eficiencia en el razonamiento matemático.

• Reinforcement Learning (RL): OLLM presenta un camino prometedor para el RL en LLMs. Al operar en un espacio de acciones de baja dimensión (las opciones latentes) en lugar de todo el vocabulario, la optimización de recompensas se vuelve mucho más eficiente en muestras y reduce la varianza.
• Interpretabilidad: La estructura de opciones ofrece interpretabilidad, donde diferentes latentes pueden corresponder a trayectorias de razonamiento distintas, permitiendo la exploración controlada de soluciones alternativas.
• Generalización: La arquitectura es lo suficientemente simple como para aplicarse a cualquier LLM preentrenado, sugiriendo su utilidad en dominios más allá de las matemáticas, como la generación de código, preguntas científicas y planificación multi-paso.

En resumen, OLLM propone un cambio de paradigma: en lugar de tratar la incertidumbre como ruido a ser muestreado, la trata como una estructura latente aprendible que puede ser controlada eficientemente por políticas externas.