Resumen Técnico: La Calibración Probabilística es una Capacidad Entrenable en Modelos de Lenguaje

Enunciado del Problema

Los modelos de lenguaje (LM) se despliegan cada vez más en escenarios que requieren aleatoriedad especificada por el usuario, como generar salidas diversas, equilibrar probabilidades de elección o muestrear distribuciones matemáticas específicas. Sin embargo, los LM autoregresivos estándar son malos generadores estocásticos nativos. Sus distribuciones de salida inducidas a menudo colapsan en subconjuntos estrechos, violan restricciones de soporte y no se alinean con las distribuciones objetivo solicitadas. Las soluciones existentes mediante prompting (por ejemplo, inyectar razonamiento o semillas) pueden mitigar parcialmente estos problemas, pero introducen costos en tiempo de inferencia, sensibilidad al prompt y sobrecarga computacional sin alterar fundamentalmente el comportamiento de muestreo del modelo.

La pregunta central abordada es si la calibración probabilística—la capacidad de un LM para generar muestras que reflejen una distribución objetivo conocida—puede aprenderse directamente mediante ajuste fino, en lugar de depender de heurísticas en tiempo de inferencia.

Metodología: Ajuste Fino para Calibración

Los autores proponen un marco para entrenar LM como muestreadores probabilísticos calibrados utilizando prompts sintéticos que solicitan explícitamente el muestreo de distribuciones matemáticas (por ejemplo, "Genera exactamente UN número aleatorio de una distribución [distribución] con parámetros [parámetros]").

Se comparan dos estrategias de ajuste fino distintas, ambas utilizando Adaptación de Bajo Rango (LoRA) en modelos base congelados:

1. Ajuste Fino con Objetivo Suave:

   • Mecanismo: La distribución objetivo se discretiza en un espacio de salida canónico finito. Se construye un trie de prefijos a partir de todas las salidas numéricas válidas, donde cada nodo representa un prefijo de token.
   • Supervisión: Para cualquier prefijo dado, el método calcula una distribución de siguiente token objetivo normalizando la masa de probabilidad restante de todas las continuaciones válidas en el trie.
   • Objetivo: El modelo minimiza la divergencia de Kullback-Leibler (KL) entre su distribución de siguiente token predicha y la distribución objetivo derivada del trie en cada prefijo visitado. Esto proporciona supervisión densa a nivel de prefijo.

2. Ajuste Fino con Objetivo Rígido:

   • Mecanismo: El modelo se entrena sobre completaciones muestreadas directamente de la distribución objetivo.
   • Supervisión: Para cada prompt, se muestrea una respuesta numérica de la distribución objetivo, se tokeniza y se utiliza como la completación de verdad fundamental.
   • Objetivo: El modelo minimiza la entropía cruzada autoregresiva enmascarada estándar sobre la completación muestreada. Esto proporciona supervisión dispersa a nivel de trayectoria, requiriendo más pasos de entrenamiento (más muestras por prompt) para lograr una cobertura comparable.

Configuración Experimental:

• Modelos: 12 modelos de cuatro familias (Qwen3, Gemma-3, Llama-3.2, GPT-OSS) que van desde 0.6B hasta 27B parámetros.
• Datos: 30 familias de distribuciones (24 vistas para entrenamiento, 6 retenidas para pruebas fuera de distribución (OOD)).
• Evaluación:
  • Muestreo Estructurado: Medido mediante la distancia Wasserstein-1 (W1) y la divergencia KL a nivel de logits contra distribuciones objetivo.
  • Benchmarks de Transferencia: Generación aleatoria abierta (soporte de diversidad), equilibrado de posiciones de respuesta en preguntas de opción múltiple (MCQ) y NoveltyBench (diversidad semántica).
  • Retención: Retención de capacidades en TinyBenchmarks (MMLU, GSM8K, etc.) y ajuste de modelado de lenguaje en PALOMA.

Contribuciones Clave

1. Entrenabilidad de la Calibración: El artículo demuestra que la calibración probabilística es una capacidad entrenable. Tanto el ajuste fino con objetivo suave como el de objetivo rígido mejoran sustancialmente la fidelidad del muestreo estructurado en familias de distribuciones retenidas y configuraciones de parámetros no vistas.
2. Análisis Comparativo de Estrategias:
   • Objetivo Rígido: Generalmente logra un rendimiento superior en el benchmark de muestreo numérico estructurado en dominio, particularmente para parámetros no vistos. Destaca en la fidelidad numérica exacta.
   • Objetivo Suave: Aunque es competitivo en tareas numéricas, demuestra una transferencia superior a configuraciones de generación estocástica más amplias. Amplía consistentemente el soporte estocástico en generación abierta, mejora el equilibrado de posiciones de respuesta en MCQ y produce una mayor utilidad en NoveltyBench.
3. Compensaciones y Costos: El estudio destaca que las ganancias en calibración a menudo conllevan compensaciones. Específicamente, el razonamiento aritmético (GSM8K) tiende a retroceder, con costos que varían según el tamaño del modelo y la variante de ajuste fino. Sin embargo, el ajuste de modelado de lenguaje (perplejidad PALOMA) a menudo se preserva o mejora, lo que sugiere que los modelos no simplemente aplanan las probabilidades de los tokens, sino que aprenden comportamientos estocásticos específicos.

Resultados

• Muestreo Estructurado: Ambos métodos reducen significativamente el W1 normalizado y el KL de logits en comparación con los puntos de control base. El ajuste fino con objetivo rígido es típicamente el mejor desempeño en benchmarks numéricos, reduciendo los errores en un orden de magnitud en muchos casos.
• Generación Abierta: El ajuste fino con objetivo suave aumenta consistentemente el tamaño de soporte del top-90% (el número de tokens necesarios para cubrir el 90% de la masa de probabilidad) en todos los modelos, a menudo en uno o dos órdenes de magnitud. Esto conduce a salidas más diversas y menos concentradas en categorías como ciudades, animales y nombres.
• MCQ y Novedad: Las configuraciones con objetivo suave generalmente producen un mejor equilibrio de posiciones de respuesta en la generación de MCQ y puntuaciones de utilidad más altas en NoveltyBench, lo que indica que el comportamiento estocástico aprendido se transfiere a restricciones implícitas del lenguaje natural.
• Retención: El impacto en las capacidades generales es mixto. Mientras que algunos modelos de tamaño mediano a grande mantienen o mejoran las puntuaciones agregadas en TinyBenchmarks, los modelos más pequeños y tareas específicas (notablemente GSM8K) muestran degradación. Esto sugiere que el objetivo de entrenamiento (completaciones cortas y directas sin trazas de razonamiento) puede entrar en conflicto con tareas pesadas en razonamiento.

Significado y Afirmaciones

El artículo afirma que la calibración probabilística no es una limitación inherente de los modelos de lenguaje, sino una capacidad que puede entrenarse explícitamente. Al plantear el problema como el aprendizaje tanto del espacio de salida válido como de la asignación de masa de probabilidad en él, los autores muestran que los LM pueden transformarse en muestreadores probabilísticos controlados.

El significado radica en ir más allá del prompting en tiempo de inferencia hacia soluciones en tiempo de entrenamiento que alteran fundamentalmente el comportamiento distribucional del modelo. Los autores señalan que, aunque las ganancias son selectivas y no mejoran universalmente todas las tareas posteriores (siendo el razonamiento aritmético una excepción notable), la capacidad de ajustar modelos para comportamientos estocásticos específicos abre nuevas vías para aplicaciones que requieren diversidad, equidad o adhesión estricta a restricciones estadísticas. El trabajo sugiere que la investigación futura debe centrarse en los mecanismos aprendidos por estos objetivos y en desarrollar variantes conscientes de la retención para mitigar la degradación de capacidades.