Resumen Técnico: N-GRPO - Mezcla de Vecinos a Nivel de Embedding para una Optimización de Política Mejorada

1. Planteamiento del Problema

El éxito de los Modelos de Lenguaje de Gran Escala (LLMs) en el razonamiento matemático depende en gran medida de la generación de trayectorias de solución diversas y válidas durante la fase de despliegue (rollout) del Aprendizaje por Refuerzo (RL). Sin embargo, las técnicas de exploración actuales enfrentan un compromiso fundamental:

• Muestreo a Nivel de Token: Los métodos estándar suelen producir trayectorias redundantes que difieren solo en el reformuleo o en reordenamientos conmutativos (por ejemplo, "1+2" frente a "2+1") sin alterar la lógica de razonamiento subyacente.
• Ruido a Nivel de Embedding: Los métodos que inyectan ruido aleatorio (por ejemplo, ruido Gaussiano) en los embeddings continuos frecuentemente interrumpen la consistencia semántica. Experimentos preliminares mostraron que el ruido sin restricciones desplaza las representaciones fuera del manifold semántico, transformando tokens válidos en otros no relacionados y descarrilando las trayectorias de despliegue.

Los intentos existentes para cerrar esta brecha, como Soft Thinking o HRPO, a menudo dependen del muestreo discreto de tokens para su aleatoriedad o no logran mantener la coherencia semántica al perturbar el espacio continuo.

2. Metodología: N-GRPO

Los autores proponen N-GRPO, una novedosa estrategia de exploración integrada en el marco de Group Relative Policy Optimization (GRPO). La innovación central es la Mezcla de Vecinos Semánticos (Semantic Neighbor Mixing), un mecanismo de perturbación a nivel de embedding que inyecta diversidad mientras se adhiere estrictamente al manifold semántico local.

Mecanismo Central: Mezcla de Vecinos Semánticos

En lugar de muestrear tokens discretos o añadir ruido sin restricciones, el método construye representaciones de entrada mezclando el embedding de un token ancla con sus vecinos semánticos más cercanos.

1. Selección del Ancla: En cada paso de generación, el modelo identifica el token más preferido (el ancla semántica, $v^*_t$) basado en los logits escalados por temperatura.
2. Recuperación de Vecinos: Se recupera un conjunto de $k$ tokens vecinos ($C_t$) de la matriz de embeddings basándose en la similitud de coseno con el ancla. Esto asegura que los vecinos se agrupen a lo largo de la misma dirección semántica.
3. Mezcla Ponderada: Los logits escalados por temperatura se normalizan sobre este conjunto de vecinos para calcular los pesos de mezcla ($\alpha_t$).
4. Construcción de Embedding Continuo: Se genera un embedding mixto continuo ($\tilde{e}_{t+1}$) como la suma ponderada de los embeddings de los vecinos. Este sirve como la entrada para el siguiente paso de tiempo.

Integración con GRPO

N-GRPO integra este mecanismo en la fase de despliegue de GRPO mediante un mecanismo de compuerta estocástica:

• Máscara de Mezcla: Una máscara de Bernoulli ($m_{i,t}$) con una probabilidad fija $\rho$ (por ejemplo, 0.1) determina si se activa la Mezcla de Vecinos Semánticos o si se sigue el muestreo de tokens discretos estándar.
• Estabilidad: Al mantener la mayoría de los pasos como tokens discretos, el método preserva la estabilidad semántica mientras introduce la intermitencia necesaria para la exploración.
• Entrenamiento: Durante la fase de entrenamiento, los conjuntos candidatos y los pesos se registran para asegurar una reconstrucción de trayectoria reproducible para el cálculo de la ventaja. La recompensa final se calcula basándose en los tokens ancla discretos correspondientes a los pasos mezclados.

3. Contribuciones Clave

• Exploración Controlable a Nivel de Embedding: El artículo introduce un método que inyecta perturbaciones guiadas semánticamente durante la generación, cerrando la brección entre el razonamiento discreto a nivel de token y la exploración continua a nivel de embedding.
• Marco N-GRPO: Los autores integran la Mezcla de Vecinos Semánticos en el marco de GRPO, creando un sistema que utiliza los manifolds semánticos locales para expandir el espacio de exploración sin sacrificar la coherencia.
• Validación Empírica: Evaluaciones extensas en diferentes escalas de modelos (1.5B a 7B) y arquitecturas (DeepSeek-R1-Distill-Qwen, Llama-3.2, Qwen3) demuestran mejoras consistentes sobre bases sólidas.

4. Resultados Experimentales

El método fue evaluado en desafíos de razonamiento matemático (AIME25, AMC23, MATH500) y tareas fuera de distribución (OOD) (GPQA-Diamond).

• Razonamiento Matemático: N-GRPO superó consistentemente a las bases, incluyendo GRPO estándar, Soft Thinking y la base STHT (que utiliza ruido Gaussiano sin restricciones).
  • En el modelo DeepSeek-R1-Distill-Qwen-1.5B, N-GRPO mejoró el promedio de Pass@32 de 77.41 (GRPO) a 79.17.
  • En el modelo 7B, el promedio de Pass@32 mejoró de 81.94 (GRPO) a 84.20.
  • Se observaron ganancias significativas en el desafiante benchmark AIME25, donde N-GRPO logró un Pass@32 de 50.28 (1.5B) y 61.18 (7B), superando a GRPO y STHT.
• Generalización Fuera de Distribución (OOD): En el benchmark GPQA-Diamond, N-GRPO demostró una generalización robusta, alcanzando las puntuaciones de Pass@32 más altas (92.87 para 1.5B y 92.80 para 7B), lo que indica que el método mejora las capacidades de razonamiento científico general, no solo las habilidades específicas de matemáticas.
• Transferibilidad: El mecanismo se transfirió con éxito al marco GSPO (N-GSPO), mostrando una ganancia de +7.66 en AIME25, lo que sugiere que el mecanismo de perturbación no está estrechamente acoplado a la normalización de ventaja específica de GRPO.
• Estudios de Ablación:
  • Tasa de Mezcla: El rendimiento se mantuvo estable a través de diversas tasas de mezcla (0.1 a 0.2), aunque una mezcla excesiva (0.2) en modelos pequeños para AIME25 causó una caída, resaltando la necesidad de un equilibrio entre diversidad y estabilidad.
  • Métrica de Distancia: La distancia de coseno superó significativamente a las métricas L1 y L2 para la selección de vecinos, debido a la naturaleza direccional de la información semántica en espacios de embedding de alta dimensión.
  • Inferencia: Los autores señalan que aplicar la Mezcla de Vecinos Semánticos durante la fase de inferencia (sin el mecanismo de filtrado de GRPO) degrada el rendimiento, confirmando su papel como herramienta de exploración en tiempo de entrenamiento en lugar de una estrategia de decodificación.

5. Significado y Reivindicaciones

El artículo afirma que N-GRPO ofrece una vía prometedora para avanzar en las capacidades de razonamiento de los LLM al explorar la dinámica del razonamiento en el espacio de los embeddings. Su importancia radica en:

• Resolver el Compromiso: Logra inyectar diversidad en el proceso de despliegue sin la deriva semántica causada por la inyección de ruido ingenuo o la redundancia del muestreo a nivel de token.
• Fundamentación Semántica: Al aprovechar el manifold semántico local, el método asegura que la exploración permanezca dentro de regiones de razonamiento válidas.
• Generalización: El enfoque no se limita a familias de modelos específicos o políticas destiladas de razonamiento, mostrando robustez a través de diferentes escalas y modelos base.

Los autores concluyen que, si bien el método introduce una carga computacional adicional durante el despliegue y actualmente se valida principalmente en tareas de matemáticas y ciencias, establece un nuevo paradigma para la optimización de políticas que equilibra la exploración y la consistencia semántica.