Resumen Técnico: Cuando la distancia distrae: Sesgo de distancia de representación en BT-Loss para modelos de recompensa

1. Planteamiento del problema

Los modelos de recompensa (RM) son fundamentales para el Aprendizaje por Refuerzo a partir de la Retroalimentación Humana (RLHF), actuando como el sustituto de las preferencias humanas que guían la alineación de los Grandes Modelos de Lenguaje (LLM). El objetivo estándar para entrenar estos modelos es la pérdida de Bradley-Terry (BT), que aprende de datos de pares consistentes en respuestas "elegidas" ($y_w$) y "rechazadas" ($y_l$).

Los autores identifican un fallo crítico en la dinámica de actualización de la pérdida BT estándar: el Sesgo de Distancia de Representación. Mientras que se pretende que el gradiente de BT sea impulsado por el error de predicción (la magnitud del error en el ranking), el análisis revela que la norma del gradiente escala con dos componentes distintos:

1. Error de Predicción: La diferencia en las recompensas predichas entre las respuestas elegida y rechazada.
2. Distancia de Representación: La distancia entre las dos respuestas en el espacio de salida de la última capa ($\|h_w - h_l\|$).

Este acoplamiento crea una señal de aprendizaje espuria donde la magnitud de la actualización está dictada por qué tan alejadas están las representaciones del par, en lugar de centrarse únicamente en qué tan equivocado está el modelo.

• Pares de pequeña distancia: Los pares que son semántica o estructuralmente similares (por ejemplo, un fragmento de código correcto frente a uno con un único error lógico) suelen tener distancias de representación pequeñas. Bajo la pérdida BT estándar, estos pares reciben actualizaciones extremadamente débiles, incluso cuando el modelo los clasifica erróneamente. Esto es problemático porque las distinciones finas son a menudo las más críticas para tareas como el razonamiento y la programación.
• Pares de gran distancia: Los pares que son muy diferentes (por ejemplo, una negativa frente a una respuesta a una consulta dañina) tienen naturalmente distancias de representación grandes. Estos pares reciben actualizaciones desproporcionadamente fuertes, lo que puede eclipsar las señales de aprendizaje de los pares más sutiles y finos.

Consecuentemente, el modelo se entrena insuficientemente en las distinciones difíciles y finas donde más falla, mientras que sobrepondera las diferencias superficiales.

2. Metodología: NormBT

Para abordar este sesgo, los autores proponen NormBT, un esquema de normalización de pares adaptativo que reescala la contribución del gradiente de cada par de preferencias.

Base Teórica

Los autores descomponen la norma del gradiente BT $\|\nabla_\theta L_{BT}\|$ y demuestran que es proporcional a:
$$\|\nabla_\theta L_{BT}\| \approx |\sigma(d) - 1| \cdot \|h_w - h_l\|$$
donde $d$ es la diferencia de recompensa y $\|h_w - h_l\|$ es la distancia de representación. Para lograr que el tamaño de la actualización dependa principalmente del término de error de predicción $|\sigma(d) - 1|$, los autores introducen un peso $w_i$ que escala inversamente con la distancia de representación.

Detalles de Implementación

Calcular directamente la norma exacta del gradiente para el pesaje es computacionalmente costoso. En su lugar, NormBT utiliza la distancia de representación de la última capa ($\|h_w - h_l\|$) como un proxy, la cual se demuestra que es un proxy exacto para la cabeza de puntuación lineal y una fuerte aproximación para el modelo completo.

La función de pérdida propuesta es:
$$\mathcal{L}_{NormBT}(\theta) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \tilde{w}(y_w, y_l) \cdot \log \sigma(r_w - r_l) \right]$$

Donde el peso $\tilde{w}$ se define como:
$$\tilde{w}(y_w, y_l) = \frac{\mu_t}{\|h_w - h_l\| + \epsilon}$$

Decisiones de diseño clave incluyen:

• Stop-Gradient: El término de distancia se calcula con un operador de parada de gradiente (stop-gradient) para evitar que el peso de normalización sea optimizado por sí mismo.
• Estabilización mediante EMA: Para evitar la deriva de escala durante el entrenamiento (ya que las magnitudes de los embeddings pueden cambiar), los autores utilizan una Media Móvil Exponencial (EMA) de la media de la distancia de representación del lote ($\mu_t$) para normalizar los pesos. Esto asegura que los pesos sean estables y adaptativos a los cambios en la distribución.
• Integración Ligera: NormBT no requiere cambios arquitectónicos, preserva la base probabilística de BT e incurre en un gasto computacional insignificante.

3. Contribuciones Clave

1. Análisis de Gradiente: El artículo proporciona un análisis teórico y empírico detallado que muestra que las normas del gradiente de BT están determinadas conjuntamente por el error de predicción y la distancia de representación, lo que conduce a señales de aprendizaje sesgadas.
2. Identificación del Sesgo: Los autores demuestran que este acoplamiento causa que los pares de pequeña distancia (que a menudo son los más críticos para tareas finas) reciban actualizaciones débiles, mientras que los pares de gran distancia dominan la señal de aprendizaje.
3. Propuesta de NormBT: Introducen NormBT, una modificación directa (drop-in) que normaliza las actualizaciones por la distancia de representación, asegurando que las señales de aprendizaje se centren en los errores de predicción.
4. Validación Empírica: Experimentos extensos a través de diferentes backbones de LLM (Gemma-2b-it, Llama-3.2-3B) y conjuntos de datos (Unified-Feedback, Skywork-Reward) muestran mejoras consistentes en el rendimiento.

4. Resultados Experimentales

Los autores evaluaron NormBT contra la línea base estándar de BT y varias variantes (BT + Margin, BT + Label Smoothing) en el benchmark RewardBench y en tareas de selección Best-of-N (BoN) descendentes.

• Rendimiento en RewardBench: NormBT superó consistentemente a la línea base de BT en todos los entornos. Las ganancias más significativas se observaron en la categoría de Razonamiento (Reasoning), donde la precisión mejoró en más de un 5% en promedio. Esto concuerda con la hipótesis de que las tareas de Razonamiento dependen de distinciones finas (distancias de representación pequeñas) que son suprimidas por la BT estándar.
• RLHF Descendente (BoN): En experimentos de selección Best-of-N, los modelos entrenados con NormBT lograron puntuaciones "gold" (calidad de las respuestas seleccionadas) más altas en comparación con las líneas base de BT.
• Estudios de Ablación:
  • NormBT superó a otras técnicas de entrenamiento como LayerNorm y Gradient Clipping, que no lograron abordar el problema del pesaje relativo entre pares.
  • El método resultó ser complementario, mejorando el rendimiento incluso cuando se aplicaba sobre otras variantes de BT (por ejemplo, NormBT + Margin).
  • Eliminar la estabilización por EMA o utilizar métricas de similitud alternativas (similitud de coseno, embeddings de promedio de pooling) resultó en una degradación del rendimiento, validando las decisiones de diseño específicas.

5. Significado y Reivindicaciones

El artículo afirma que NormBT ofrece una solución fundamentada y ligera a una limitación estructural en la pérdida BT estándar. Al desacoplar la magnitud de la actualización de la distancia de representación, NormBT asegura que el modelo de recompensa aprenda de manera más fiel de las distinciones finas, que suelen ser las más desafiantes pero críticas para una alineación de alta calidad.

Los autores enfatizan que su trabajo no aborda directamente la calidad de los datos o el sesgo del anotador, sino que mejora la eficiencia y fidelidad de las señales de aprendizaje extraídas de los datos de pares existentes. Posicionan a NormBT como una mejora fundamental para el paradigma de modelado de recompensa discriminativo, lo que potencialmente conduce a señales de recompensa más fiables en aplicaciones de RLHF descendentes, particularmente en dominios que requieren razonamiento preciso y detección de errores sutiles.