Resumen Técnico: LARK (Selección de Trayectorias Basada en la Aprendibilidad)

Planteamiento del Problema

El artículo aborda el desafío de la destilación de razonamiento, un paradigma donde un modelo de lenguaje (LLM) estudiante más pequeño es ajustado (fine-tuned) para imitar las trayectorias de razonamiento generadas por modelos maestros más grandes. Mientras que los métodos actuales de selección de datos suelen operar a nivel de pregunta o dependen de heurísticas como la calidad de la trayectoria, la corrección o la confianza del modelo, frecuentemente pasan por alto un factor crítico: la aprendibilidad (learnability).

Específicamente, para una sola pregunta, pueden existir múltiples trayectorias de razonamiento candidatas (generadas por diferentes maestros, ejecuciones de muestreo o estilos). Incluso si estas trayectorias conducen al mismo resultado final correcto, proporcionan señales de entrenamiento muy distintas. Una trayectoria que parece "de alta calidad" o "natural" no es necesariamente aquella de la cual el modelo estudiante puede aprender de la manera más eficiente. La pregunta central planteada es: ¿Podemos diseñar un criterio de selección de trayectorias basado en principios que garantice que las trayectorias de razonamiento destiladas sean aprendibles por el modelo estudiante?

Metodología: LARK

Los autores proponen LARK (Rankeo de Anclaje por Tiempo basado en la Aprendibilidad), un método que selecciona trayectorias basándose en una definición de aprendibilidad desde la teoría de la optimización, en lugar de métricas de calidad estáticas.

1. La Aprendibilidad como un Objetivo de Optimización

LARK enmarca la selección de trayectorias como un problema de optimización de políticas. Define un factor de aprendibilidad $\rho(\theta; q)$, que caracteriza la tasa a la que disminuye la pérdida de entrenamiento del estudiante. Este factor se define como la razón entre la norma del gradiente al cuadrado y la pérdida:
$$\rho(\theta; q) \triangleq \frac{\|\nabla_\theta \mathcal{L}(\theta; q)\|^2}{\mathcal{L}(\theta; q)}$$
donde $q$ es el vector de pesos sobre $K$ trayectorias candidatas. Bajo la condición de Polyak–Łojasiewicz (PL), un $\rho$ más alto implica una convergencia más rápida.

2. Aproximación de Tiempo de Anclaje (Anchor-Time)

Optimizar directamente el objetivo de la aprendibilidad integrado en el tiempo sobre toda la trayectoria de ajuste fino es intratable. LARK introduce la condición de anclaje relativo, que postula que, para modelos sobreparametrizados en el régimen de "entrenamiento perezoso" (lazy training), la tasa de aprendibilidad permanece proporcional a su valor en los parámetros de referencia iniciales $\theta_{ref}$. Esto permite que el objetivo sea aproximado maximizando $\rho(\theta_{ref}; q)$.

3. Linealización Local y Proxy de Pasada hacia Adelante (Forward-Pass)

Maximizar $\rho(\theta_{ref}; q)$ directamente es problemático porque es una función cuasi-convexa que colapsa en una solución "one-hot" degenerada (seleccionando una sola trayectoria). Para evitar esto, LARK:

• Linealización Local: Expande $\rho$ alrededor de la distribución uniforme $p$ (SFT estándar sin selección) utilizando una expansión de Taylor de primer orden.
• Proxy de Pasada hacia Adelante ($\hat{g}_k$): En lugar de calcular gradientes exactos (que requieren pasadas hacia atrás para cada candidato), LARK deriva un proxy computable $\hat{g}_k$ utilizando únicamente pasadas hacia adelante. Este proxy estima el gradiente local de la tasa de aprendibilidad basándose en los residuales de predicción de tokens y la pérdida de entropía cruzada del estudiante.
  • Específicamente, $\hat{g}_k$ favorece trayectorias donde el estudiante está "confundidamente equivocado" de manera estructurada: la pérdida $\ell_k$ es alta (indicando una mala predicción), pero la masa residual se concentra en unos pocos tokens de "casi acierto" (proporcionando una señal de corrección clara).

4. Selección con Regularización $\chi^2$

Para evitar que el selector "hackee" el criterio de aprendibilidad desviándose demasiado de la distribución original de los datos (lo que perjudica la generalización), LARK impone un término de regularización $\chi^2$. Esto conduce a una regla de selección de forma cerrada y parametrizada por presupuesto:

• Las trayectorias se clasifican por su puntuación $\hat{g}_k$.
• Se seleccionan las $B$ mejores trayectorias.
• A diferencia de la truncación estricta de los mejores $B$, LARK asigna pesos suaves proporcionales al margen de puntuación por encima de la $(B+1)$-ésima trayectoria. Esta regla de ponderación se deriva para maximizar el límite inferior de la mejora de la aprendibilidad mientras se controla la divergencia $\chi^2$ respecto al prior uniforme.

Principales Contribuciones

1. Formulación: Los autores formulan la selección de trayectorias de razonamiento como un problema de optimización de política basado en la aprendibilidad, introduciendo la tasa de aprendibilidad de tiempo de anclaje $\rho$ como un objetivo fundamentado.
2. Algoritmo: Derivan una regla de selección de forma cerrada y parametrizada por presupuesto que utiliza un proxy de pasada hacia adelante $\hat{g}_k$ y regularización $\chi^2$. Esto evita costosas pasadas hacia atrás por trayectoria, mientras garantiza teóricamente límites sobre el error de estimación y la generalización.
3. Validación Empírica: El artículo demuestra que LARK supera consistentemente a las líneas base existentes (incluyendo métodos basados en heurísticas, calidad y puntuación del lado del estudiante) en múltiples modelos base (Qwen-2.5-7B/1.5B, Llama-3.2-3B) y benchmarks de razonamiento (AIME, AMC, GPQA, MATH-500).

Resultados Experimentales

• Rendimiento: LARK alcanzó el promedio de Acc@5 más alto en todos los entornos de presupuesto de modelo probados. Por ejemplo, en Qwen-2.5-7B con un presupuesto de $B=3$, LARK mejoró el rendimiento en 7.07 puntos porcentuales sobre la línea base más fuerte (RSR).
• Eficiencia: LARK es un método de solo pasada hacia adelante, igualando el costo computacional de las líneas base de menor costo del lado del estudiante, mientras que los métodos exactos basados en gradientes requerirían mucha más computación.
• Análisis Diagnóstico:
  • Decaimiento de la Pérdida: Las trayectorias seleccionadas por LARK indujeron la reducción más rápida de la pérdida de SFT durante el entrenamiento.
  • Validez Predictiva: La puntuación LARK $\hat{g}_k$ mostró una fuerte correlación positiva (Pearson $r=0.69$) con la precisión downstream cuando se agregó por modelo maestro.
  • Comportamiento a Nivel de Token: Las trayectorias seleccionadas exhibieron probabilidades de tokens correctos más bajas pero una mayor "concentración de tokens erróneos", confirmando que LARK selecciona ejemplos donde el estudiante está equivocado pero el error es estructurado y corregible.

Significado y Reivindicaciones

El artículo afirma que la aprendibilidad es un principio distinto y efectivo para la destilación de razonamiento eficiente, separado de la calidad o dificultad de la trayectoria. Al fundamentar la selección en la dinámica de optimización del modelo estudiante, LARK identifica trayectorias que proporcionan la señal de aprendizaje más eficiente.

Los autores enfatizan que su método:

• Preserva la Generalización: A través de la regularización $\chi^2$, evita que la política de selección se sobreajuste a un conjunto estrecho de trayectorias "fáciles" o "hackeadas".
• Está Fundamentado Teóricamente: La regla de selección se deriva de los primeros principios de la teoría de la optimización (condiciones PL, flujo de gradiente) con límites de error explícitos.
• Es Práctico: No requiere ajuste de hiperparámetros para la regla de selección en sí misma (la fuerza de la regularización está determinada por el presupuesto $B$) y evita la sobrecarga computacional de las pasadas hacia atrás.

El trabajo sugiere que futuras extensiones podrían aplicar esta perspectiva basada en la aprendibilidad a conjuntos de candidatos más ruidosos y dominios de razonamiento más amplios, pero el alcance actual se limita estrictamente a la destilación de razonamiento en la resolución de problemas matemáticos.