Learning Adaptive LLM Decoding

Este trabajo propone el uso de adaptadores de decodificación ligeros entrenados con aprendizaje por refuerzo para seleccionar dinámicamente estrategias de muestreo adaptativas a nivel de secuencia y token, logrando mejoras significativas en la precisión de tareas matemáticas y de codificación bajo restricciones de recursos computacionales fijos.

Lo esencial

Imagina que tienes un genio muy inteligente (un modelo de lenguaje grande o LLM) que vive dentro de una caja. Este genio puede resolver problemas de matemáticas, escribir código o contar historias. Pero hay un problema: para que el genio hable, tú tienes que darle instrucciones sobre cómo debe hablar.

Hasta ahora, esas instrucciones eran fijas y aburridas. Era como decirle al genio: "Siempre habla con un tono de voz medio, eligiendo palabras al azar de una lista de 10 opciones". No importaba si el problema era fácil (como "2+2") o difícil (como "resolver un misterio de espionaje"). El genio usaba el mismo tono y la misma estrategia para todo. A veces, para problemas fáciles, esto era un desperdicio de energía. Otras veces, para problemas difíciles, el genio se atascaba porque no se le permitía "pensar en voz alta" con suficiente libertad.

La Gran Idea: El "Adaptador de Aprendizaje"

Los autores de este paper proponen algo revolucionario: no cambiar al genio, sino darle un "asistente de decisiones" (un adaptador) que aprenda a elegir la mejor forma de hablar en cada momento.

Piensa en este adaptador como un director de orquesta o un entrenador deportivo que se sienta junto al genio. Su trabajo no es escribir la respuesta, sino decidir cómo el genio debe generar cada palabra.

El sistema funciona en dos niveles, como si tuvieras dos tipos de decisiones:

1. El Nivel de la "Historia Completa" (Adaptador a nivel de secuencia)

Imagina que le das al genio un problema nuevo. Antes de que empiece a escribir la primera palabra, el entrenador mira el problema y dice:

• "Este es un problema de matemáticas muy difícil. ¡Necesitamos que el genio sea creativo y explore muchas posibilidades! Vamos a usar un tono de voz más libre y aleatorio."
• "Este es un problema de código simple. ¡No hace falta arriesgarse! Vamos a ser muy precisos y directos."

El entrenador elige una estrategia global para toda la respuesta. Es como elegir si vas a conducir un coche por una autopista (estrategia segura y directa) o por un sendero de montaña lleno de curvas (estrategia exploradora y arriesgada).

2. El Nivel de la "Palabra por Palabra" (Adaptador a nivel de token)

Aquí es donde la magia se vuelve aún más fina. El entrenador no solo decide la estrategia al principio, sino que cambia de opinión en cada palabra que el genio escribe.

Imagina que el genio está resolviendo un problema de lógica:

• Palabra 1: "El" (Fácil, seguro). El entrenador dice: "Habla rápido y seguro, no gastes energía aquí."
• Palabra 50: "...pero si el tren viaja a..." (¡Aquí viene la parte difícil!). El entrenador nota que el genio está dudando. "¡Alto! Aquí necesitamos explorar. Cambia a un modo más creativo, prueba varias posibilidades, no te cases con una sola idea."
• Palabra 100: "...por lo tanto, la respuesta es 42." (Ya casi terminamos). El entrenador dice: "¡Vuelve a la precisión! Cierra el argumento con firmeza."

El entrenador decide en tiempo real: ¿Debo ser arriesgado ahora o debo ser conservador? Todo esto depende de cuánto "presupuesto" (energía de computación) le queda al genio para seguir pensando.

¿Cómo aprende el entrenador? (El Entrenamiento)

Lo más genial es que no les dicen al entrenador qué hacer. No hay un manual de instrucciones escrito por humanos.

En su lugar, usan un sistema de premios y castigos (como en los videojuegos o el entrenamiento de perros):

1. El entrenador prueba una estrategia.
2. El genio genera una respuesta.
3. Si la respuesta es correcta (resuelve el problema de matemáticas o el código funciona), el entrenador recibe un premio.
4. Si falla, recibe un castigo.

Con el tiempo, el entrenador aprende por prueba y error: "¡Oh! Cuando el problema es difícil y me queda mucho tiempo, si elijo la estrategia 'exploradora', gano más premios. Pero si me queda poco tiempo, es mejor ser 'conservador'."

¿Por qué es importante esto?

• Ahorro de energía: No gastas tiempo de computación explorando opciones en problemas fáciles.
• Mejor calidad: Cuando el problema es difícil, el sistema sabe cuándo "soltar la rienda" y dejar que el genio explore ideas locas que podrían llevar a la solución correcta.
• Flexibilidad: El mismo genio puede resolver problemas de matemáticas y de programación, y el entrenador sabe cambiar el estilo de "conducción" según sea necesario.

En resumen

Este paper nos dice que la forma en que un modelo de IA "piensa" (decodifica) es tan importante como lo que "sabe". En lugar de tener un interruptor fijo para todo, hemos creado un piloto automático inteligente que ajusta la velocidad, la dirección y el riesgo en cada segundo del viaje, asegurándose de llegar a la meta (la respuesta correcta) de la manera más eficiente posible.

Es como pasar de conducir un coche con el freno de mano puesto y el acelerador fijo, a tener un copiloto experto que sabe exactamente cuándo acelerar, cuándo frenar y cuándo tomar un atajo, dependiendo del terreno y del combustible que te queda.

Resumen técnico

Resumen Técnico: Aprendizaje de Decodificación Adaptativa para LLMs

1. Planteamiento del Problema

La inferencia de Modelos de Lenguaje Grande (LLMs) suele depender de hiperparámetros de muestreo fijos (como temperatura, top-k, top-p o min-p) que se seleccionan estáticamente para todo un modelo o conjunto de datos. Esta aproximación ignora la heterogeneidad significativa existente entre diferentes prompts, estilos de razonamiento e incluso entre tokens individuales dentro de una misma secuencia.

• Ineficiencia: Las estrategias fijas no se adaptan a la incertidumbre latente del modelo. Por ejemplo, ciertos tokens críticos en tareas de razonamiento (conocidos como "tokens de bifurcación" o forking tokens) requieren mayor exploración estocástica, mientras que otros pueden resolverse de manera determinista.
• Desconexión Entrenamiento-Inferencia: En los marcos actuales de Aprendizaje por Refuerzo con Recompensas Verificables (RLVR), la estrategia de decodificación se mantiene fija durante la generación, creando una desconexión entre la optimización del modelo y las restricciones de cómputo reales en tiempo de inferencia.
• Objetivo: El trabajo propone aprender políticas de decodificación adaptativas que seleccionen dinámicamente la estrategia de muestreo en tiempo de inferencia, condicionadas a los recursos de cómputo disponibles, sin necesidad de ajustar los parámetros del modelo de lenguaje base.

2. Metodología

Los autores introducen los Adaptadores de Decodificación Aprendidos (Learned Decoding Adapters), una familia de políticas ligeras entrenadas con Aprendizaje por Refuerzo (RL) que modulan el comportamiento de muestreo de un LLM congelado.

El enfoque se divide en dos niveles de granularidad:

A. Nivel de Secuencia (Contextual Bandit)

• Formulación: Se modela como un problema de bandito contextual. Para cada prompt, el adaptador selecciona una única configuración de decodificación (ej. greedy, top-k, top-p) que se aplica uniformemente a toda la generación.
• Entrada: El vector de características incluye la representación del prompt (embedding) y el presupuesto de muestreo paralelo ($B$), que define cuántas trayectorias completas se pueden generar.
• Acción: Selecciona un par de hiperparámetros de un espacio de acciones discreto predefinido.
• Entrenamiento: Se utiliza el algoritmo REINFORCE (gradiente de política) para maximizar la recompensa terminal (ej. corrección en matemáticas o código) bajo el presupuesto dado.

B. Nivel de Token (Proceso de Decisión de Markov Parcialmente Observable - POMDP)

• Formulación: Se modela como un POMDP donde el adaptador toma decisiones en cada paso de generación.
• Entrada: Observa las representaciones internas del modelo (embeddings de estados ocultos) y el presupuesto de tokens restantes ($b_t$).
• Acción: Selecciona una acción de decodificación en cada paso. En los experimentos, el espacio de acciones se restringió a variaciones de temperatura, ya que se identificó como un eje interpretable y efectivo para controlar la estocasticidad dentro de una sola trayectoria.
• Estabilidad: Para evitar gradientes de alta varianza, se filtran prompts con señales de recompensa ruidosas y se enmascaran tokens con distribuciones de probabilidad altamente concentradas (probabilidad máxima > 0.95).

C. Selección del Espacio de Acciones

• Se utiliza un procedimiento de selección basado en la cobertura (inspirado en AuPair) para elegir un subconjunto pequeño y diverso de estrategias de decodificación a partir de un gran pool candidato. Esto asegura que el espacio de acciones cubra comportamientos cualitativamente distintos sin redundancia.

3. Contribuciones Clave

1. Formulación Unificada: Plantean la inferencia en tiempo de decodificación como un problema de aprendizaje de políticas, introduciendo un marco de RL unificado para la adaptación a nivel de prompt y de token bajo presupuestos de cómputo explícitos.
2. Entrenamiento sin Modelos de Recompensa: Los adaptadores se entrenan exclusivamente con recompensas verificables de la tarea (corrección de código/matemáticas), sin necesidad de modelos de recompensa aprendidos, etiquetas de preferencia o heurísticas de diseño manual.
3. Eficiencia y Adaptabilidad: Demuestran que es posible mejorar el rendimiento de razonamiento manteniendo el modelo base congelado, aprendiendo a asignar estocasticidad de manera óptima según la dificultad del problema y los recursos disponibles.

4. Resultados Experimentales

Los experimentos se realizaron en los conjuntos de datos MATH (resolución de problemas matemáticos) y CodeContests (programación competitiva), utilizando modelos como Qwen3-4B y Qwen2.5-Math.

• Rendimiento General: Los adaptadores aprendidos superaron consistentemente a las estrategias estáticas de referencia (la mejor estrategia fija y mezclas fijas).
• Nivel de Token: El adaptador a nivel de token mejoró la precisión Pass@1 en un 10.2% sobre la mejor línea base estática en el conjunto MATH bajo un presupuesto de tokens fijo.
• Nivel de Secuencia: El adaptador a nivel de secuencia mostró ganancias de 2-3% en escenarios con muestreo paralelo limitado.
• Condicionamiento al Presupuesto: Entrenar la política condicionada al presupuesto de cómputo (número de rollouts o tokens restantes) mejoró significativamente el rendimiento en comparación con entrenamientos agnósticos al presupuesto.
• Generalización: Los adaptadores entrenados en un dominio (ej. Matemáticas) mostraron capacidad de generalización a dominios no vistos (ej. Código o problemas más difíciles como AIME 2025), aunque con ganancias menores que en el dominio de entrenamiento.
• Análisis de Comportamiento: Se observó que la política aprendida tiende a colapsar tokens de baja entropía a comportamientos casi deterministas, mientras que mantiene la estocasticidad en tokens de alta entropía (puntos críticos), aunque no sigue reglas heurísticas simples como umbrales de entropía fijos.

5. Significado e Impacto

Este trabajo identifica el control en tiempo de inferencia como un eje fundamental y subexplorado para mejorar el razonamiento de los LLMs, complementario al escalado de modelos y al ajuste fino (fine-tuning).

• Eficiencia Computacional: Permite obtener mejores resultados sin aumentar el tamaño del modelo, optimizando el uso de recursos de cómputo disponibles.
• Flexibilidad: Ofrece un mecanismo para adaptar dinámicamente la exploración vs. explotación durante la generación, crucial para tareas complejas de razonamiento donde la incertidumbre varía a lo largo de la secuencia.
• Paradigma de Aprendizaje: Propone un cambio de paradigma donde la estrategia de decodificación no es un hiperparámetro estático, sino una política aprendible que interactúa con el entorno del modelo congelado.

En conclusión, el artículo demuestra que aprender a "cómo muestrear" es tan importante como "qué modelo usar", logrando mejoras sustanciales en tareas de razonamiento mediante políticas ligeras y adaptativas.