Resumen Técnico: F-TIS – Aprovechamiento de Modelos Diversos en GRPO Colaborativo

Declaración del Problema

La Optimización de Política Relativa de Grupo (GRPO) ha surgido como un método dominante para el entrenamiento posterior de Modelos de Lenguaje Grandes (LLM), particularmente para mejorar las capacidades de razonamiento. A diferencia de la Optimización de Política Próxima (PPO), GRPO elimina la necesidad de un modelo de valor separado al calcular las ventajas relativas a un grupo de completaciones generadas para un único prompt. Sin embargo, la naturaleza auto-regresiva de la generación de LLM crea un cuello de botella computacional significativo durante la fase de inferencia.

Para abordar esto, trabajos anteriores han explorado distribuir el paso de generación a través de múltiples nodos. Los enfoques distribuidos existentes asumen en gran medida modelos homogéneos (arquitecturas y parámetros idénticos) para mantener condiciones "en política". Esta suposición es impráctica para escenarios de entrenamiento descentralizado donde los participantes poseen recursos y preferencias heterogéneos. En tales entornos, surgen tres tipos de heterogeneidad:

1. Tamaño del Modelo: Los nodos pueden ejecutar modelos con conteos de parámetros muy diferentes (por ejemplo, 1B vs. 3B).
2. Experiencia del Modelo: Los nodos pueden poseer modelos de tamaño similar pero con especializaciones diferentes (por ejemplo, generalista vs. programador).
3. Parámetros Entrenables: Los nodos pueden emplear diferentes estrategias de ajuste fino (por ejemplo, ajuste fino completo vs. Ajuste Fino Eficiente en Parámetros como LoRA).

Cuando modelos heterogéneos colaboran sin ajuste, los datos de entrenamiento se vuelven altamente fuera de política. La divergencia entre la política generadora ($\pi_{gen}$) y la política de actualización ($\pi_{\theta}$) introduce ruido destructivo, lo que lleva al colapso de la política o a un rendimiento significativamente degradado en comparación con el entrenamiento en aislamiento.

Metodología: Muestreo de Importancia Filtrado y Recortado (F-TIS)

Los autores proponen F-TIS, un marco unificado diseñado para permitir un entrenamiento colaborativo GRPO estable a través de modelos heterogéneos mientras se minimiza la sobrecarga de comunicación. La metodología combina dos mecanismos clave: Muestreo de Importancia Recortado (TIS) y Filtrado de Muestras.

1. Muestreo de Importancia Recortado (TIS)

GRPO estándar asume datos en política. Para manejar muestras fuera de política, los autores adaptan la función de pérdida utilizando muestreo de importancia. Comparan tres enfoques:

• NoIS (Sin Muestreo de Importancia): Trata todas las muestras como en política (línea base para el fallo).
• VIS (Muestreo de Importancia Vanilla): Aplica la razón de importancia $\frac{\pi_{\theta}}{\pi_{gen}}$ a cada token.
• TIS (Muestreo de Importancia Recortado): Recorta la razón de importancia en una constante $C$ (por ejemplo, $C=2$) para evitar que pesos extremos desestabilicen el entrenamiento. La razón se mueve fuera de la suma interna sobre los tokens.

La formulación de la pérdida TIS es:
$$\mathcal{L}_{GRPO} = \frac{1}{G} \sum_{i=1}^{G} \frac{1}{|a_i|} \sum_{t=1}^{|a_i|} \min \left( \frac{\pi_{\theta}(a_{i,t}|p \circ a_{i,<t})}{\pi_{gen}(a_{i,t}|p \circ a_{i,<t})}, C \right) \cdot \min \left( R_{i,\theta} \hat{A}_i, \text{clip}(R_{i,\theta} \hat{A}_i, 1-\epsilon, 1+\epsilon) \right)$$
Donde $R_{i,\theta}$ es la razón desconectada y $\hat{A}_i$ es la ventaja del grupo.

2. Filtrado de Muestras

Los autores observan que las muestras fuera de política con ventajas negativas ($\hat{A}_i < 0$) pueden amplificar tokens poco probables, lo que lleva a completaciones sin sentido. Introducen un mecanismo de filtrado basado en:

• Ventaja: Las muestras con $\hat{A}_i < 0$ se descartan para el paso de actualización.
• Divergencia KL: Las muestras donde la divergencia KL entre la política actual y la política generadora supera un umbral $g$ ($D_{KL}(\pi_{\theta} \parallel \pi_{gen}) > g$) también se filtran.

Las muestras filtradas aún se utilizan para calcular la ventaja del grupo (para mantener la clasificación relativa del grupo) pero se excluyen de la actualización del gradiente.

3. El Marco F-TIS

F-TIS integra TIS y filtrado. El término de ventaja $\hat{A}_{t,i}$ en la función de pérdida se establece en 0 si la muestra se filtra (es decir, si $\hat{A}_i \leq 0$ O $D_{KL} > g$).

• Eficiencia de Comunicación: El marco requiere comunicación mínima, transmitiendo solo tokens y log-probabilidades (aprox. 8 bytes por token), lo que lo hace adecuado para redes descentralizadas.

Contribuciones Clave

1. Identificación de Desafíos de Heterogeneidad: El documento demuestra empíricamente que GRPO estándar falla en entornos descentralizados heterogéneos (NoIS), lo que lleva a una degradación del rendimiento en comparación con el entrenamiento en aislamiento.
2. Algoritmo F-TIS: Un nuevo paradigma de entrenamiento que aprovecha con éxito muestras fuera de política de modelos diversos (diferentes tamaños, experiencia y configuraciones PEFT) sin colapsar la política.
3. Eficiencia de Comunicación: El enfoque mantiene una sobrecarga de comunicación baja, requiriendo solo log-probabilidades y tokens, a diferencia de métodos que requieren sincronización completa del modelo o intercambios pesados de modelos de valor.
4. Estudios de Ablación: Los autores analizan sistemáticamente el impacto del umbral de filtrado $g$ y comparan F-TIS con F-VIS (Muestreo de Importancia Vanilla Filtrado), demostrando la superioridad del recorte en las etapas posteriores del entrenamiento.

Resultados Experimentales

Los autores evaluaron F-TIS en el conjunto de datos GSM8k para entrenamiento y en el conjunto de datos MATH-500 para pruebas fuera de distribución (OOD), utilizando modelos Qwen2.5 (variantes Base y Coder de 1.5B y 3B).

• Convergencia: F-TIS logra una convergencia final del modelo idéntica al entrenamiento puramente en política (homogéneo) en todos los configuraciones heterogéneas.
• Generalización (OOD): En varias configuraciones, F-TIS superó al entrenamiento en política aislado en tareas OOD. Notablemente:
  • El modelo 3B Coder emparejado con un modelo Base mejoró su rendimiento en MATH-500 en un 5.2%.
  • El modelo 3B Coder emparejado con un modelo Coder más pequeño mejoró en un 12%.
  • El modelo 3B PEFT mejoró en un 7% cuando se entrenó de manera colaborativa.
• Umbral de Filtrado ($g$): Los estudios de ablación revelaron que valores de $g$ más pequeños benefician inicialmente a los modelos más pequeños (1.5B) al eliminar ruido aleatorio, mientras que los modelos más grandes (3B) se benefician de un $g$ más alto (por ejemplo, 50) más adelante en el entrenamiento, permitiendo una mayor exploración del espacio.
• Vertical vs. Horizontal: El documento encuentra que la colaboración vertical (un nodo genera el grupo completo para un prompt) es superior a la colaboración horizontal (los nodos generan subconjuntos). La colaboración horizontal introduce sesgo en el cálculo de la ventaja relativo a la media del enjambre, degradando el rendimiento.

Significado y Afirmaciones

El documento posiciona a F-TIS como un paso fundamental hacia la RL colaborativa descentralizada para LLM. Los autores afirman que su trabajo está entre los primeros en estudiar y habilitar exitosamente la RL colaborativa a través de modelos heterogéneos.

Las afirmaciones clave incluyen:

• Viabilidad: Es posible entrenar modelos diversos juntos en una sola ejecución de GRPO sin el colapso de rendimiento típicamente asociado con datos fuera de política.
• Eficiencia: El método es eficiente en comunicación, lo que lo hace viable para sistemas descentralizados con diversas restricciones de red.
• Mejora del Razonamiento: Más allá de igualar el rendimiento en política, el uso colaborativo de muestras fuera de política puede actuar como un regularizador o impulsor de exploración, lo que lleva a una mejor generalización en tareas fuera de distribución (hasta un 12% de mejora) en comparación con entrenar modelos en aislamiento.
• Practicidad: El marco acomoda restricciones del mundo real donde los usuarios tienen diferentes capacidades de hardware (tamaños de modelo) y necesidades especializadas (modelos expertos o PEFT), permitiéndoles contribuir a un objetivo de aprendizaje compartido sin requerir infraestructura idéntica.