Multi-Task GRPO: Reliable LLM Reasoning Across Tasks

Imagina que estás entrenando a un asistente robótico superinteligente para resolver todo tipo de acertijos, desde problemas matemáticos hasta enigmas de lógica. Quieres que este robot sea un verdadero generalista, capaz de manejar cualquier cosa que le lances sin confundirse o fallar en lo difícil. En el mundo de la inteligencia artificial, esto se llama "razonamiento", y los robots son llamados Modelos de Lenguaje Extensos (LLM). Para hacer que estos modelos sean mejores, los científicos utilizan una técnica llamada "Aprendizaje por Refuerzo", que es como darle al robot un "choca esos cinco" (una recompensa) cuando acierta una respuesta y un suave "inténtalo de nuevo" cuando se equivoca. Una forma popular de hacer esto es GRPO, un método ingenioso para determinar exactamente cuánto debe aprender el robot de cada intento. Sin embargo, hay un inconveniente: si le pides al robot que aprenda diez tipos diferentes de acertijos a la vez, a menudo se queda estancado. Podría volverse muy bueno en los acertijos fáciles pero ignorar por completo los difíciles, o podría quedarse atrapado en un tipo específico de enigma mientras olvida cómo hacer los otros. La gran pregunta es: ¿cómo entrenamos a un solo cerebro para que sea igualmente bueno en todo, asegurando que ninguna habilidad se quede atrás?

Este es exactamente el problema que aborda un nuevo método llamado Multi-Task GRPO (MT-GRPO). Piensa en el proceso de entrenamiento estándar como un profesor intentando enseñar a una clase de estudiantes con velocidades de aprendizaje muy diferentes. Si el profesor solo se enfoca en el estudiante "promedio", los niños listos podrían aburrirse mientras que los que tienen dificultades se quedan cada vez más atrás. El enfoque estándar a menudo permite que las tareas fáciles "secuestren" el entrenamiento, haciendo que el modelo parezca bueno en promedio mientras secretamente falla en los desafíos más difíciles. Los autores de este artículo se dieron cuenta de que simplemente promediar las puntuaciones no crea un robot confiable; necesitas asegurarte de que el peor estudiante de la clase también esté mejorando.

Para solucionar esto, los investigadores inventaron una estrategia de dos partes. Primero, crearon un "sistema de ponderación dinámica". Imagina a un entrenador que observa constantemente el marcador. Si el robot lo está haciendo de maravilla en "Countdown" (un juego matemático) pero terriblemente en "Zebra puzzles" (acertijos de lógica), el entrenador cambia inmediatamente el enfoque. Comienzan a darle al robot más problemas de práctica de la categoría Zebra y menos de la categoría de matemáticas. Pero aquí está la parte truculenta: a veces el robot intenta un acertijo Zebra y obtiene una puntuación de "cero" porque ni siquiera intentó resolverlo adecuadamente, lo que significa que no aprende nada de ese intento. Si el entrenador solo cuenta el número de problemas, podría pensar que está enseñando acertijos Zebra, pero en realidad, el robot solo está aprendiendo de los de matemáticas porque los intentos de Zebra fueron "gradientes cero" (aprendizaje vacío).

Para resolver esto, la segunda parte de su invento es un "Muestreador de Preservación de Proporciones" (Ratio-Preserving Sampler). Esto es como un bibliotecario muy estricto que asegura que, incluso si el robot desecha los intentos de práctica inútiles, la pila final de problemas de práctica útiles todavía tenga la mezcla exacta de acertijos matemáticos y de lógica. Si el entrenador quiere que el 30% de la práctica sean acertijos de lógica, el bibliotecario se asegura de que, después de filtrar la basura, el 30% de los problemas restantes que son buenos sigan siendo acertijos de lógica. Esto evita que las tareas fáciles tomen accidentalmente el control de la sesión de entrenamiento solo porque las tareas difíciles tienen más intentos "fallidos".

Los resultados de este nuevo método son bastante prometedores. Cuando el equipo lo probó en una mezcla de tres tareas de razonamiento diferentes, su nuevo método MT-GRPO mejoró la precisión de la tarea con peor desempeño entre un 16 y 28% en comparación con el método estándar, y un 6% en comparación con otro método avanzado llamado DAPO. Lo que es más impresionante, alcanzó un umbral de precisión del 50% en las tareas más difíciles utilizando un 50% menos de pasos de entrenamiento que la competencia. En una prueba más grande con nueve tareas diferentes, el método continuó superando a otros, demostrando que puede escalar para manejar desafíos más complejos y diversos. Los autores descubrieron que, al ajustar una "perilla" (un parámetro llamado $\lambda$), podían controlar cuánto prioriza el modelo el corregir sus habilidades más débiles frente a simplemente mejorar en promedio.

En resumen, este artículo sugiere que, al ser más inteligentes sobre qué problemas practica el robot y asegurar que esos problemas realmente lleguen a la sesión de entrenamiento, podemos construir una IA que no solo sea buena en lo fácil, sino que sea confiablemente competente en todo el especto. Es un paso hacia la creación de una IA que no solo se especializa en un campo estrecho, sino que puede razonar a través de una amplia variedad de problemas del mundo real sin dejar ninguna de sus habilidades atrás.