AEGIS: Anchor-Enforced Gradient Isolation for Knowledge-Preserving Vision-Language-Action Fine-Tuning

Resumen Técnico: AEGIS

1. El Problema: Asimetría de Gradientes Cross-Modal

El artículo aborda un desafío fundamental en la adaptación de Modelos de Visión-Lenguaje (VLM) pre-entrenados para el control robótico (Modelos Visión-Lenguaje-Acción o VLA).

• El Conflicto Geométrico: Los VLMs se pre-entrenan con pérdida de Entropía Cruzada (CE), lo que moldea un paisaje de parámetros de alta dimensión y baja rank (manifold semántico complejo). Por el contrario, la fine-tuning para robótica utiliza expertos de Flow-Matching con pérdida MSE (Error Cuadrático Medio), que genera gradientes de alta magnitud pero de baja rank (concentrados en un subespacio físico estrecho de 7 grados de libertad).
• La Consecuencia: Cuando los gradientes de MSE se propagan hacia atrás en la red, sobrescriben destructivamente las representaciones semánticas delicadas entrenadas con CE. Esto provoca un olvido catastrófico de las capacidades de Respuesta a Preguntas Visuales (VQA) en muy pocos pasos de entrenamiento.
• Limitaciones de las Soluciones Actuales:
  • Stop-Gradient: Corta el flujo de gradiente continuo, preservando el VQA pero descartando la supervisión continua rica, obligando a usar tokens discretos (proxy) que comprimen mal las acciones.
  • LoRA (Low-Rank Adaptation): Restringe la capacidad de actualización pero no la geometría de la dirección del gradiente. Si el gradiente destructivo se proyecta dentro del subespacio de baja rank, LoRA lo ejecutará fielmente, causando aún así erosión del VQA.
  • Co-entrenamiento: Requiere mezclar datos de VQA en los lotes de entrenamiento, duplicando el costo computacional y requiriendo acceso a corpus originales.

2. Metodología: AEGIS (Sistema de Aislamiento de Gradientes Forzado por Anclaje)

AEGIS es un marco de intervención en la fase de retropropagación (backward-pass) que no requiere datos de replay ni cambios en la arquitectura o función de pérdida. Se basa en tres componentes clave:

1. Ancla Estática de Wasserstein (Static Wasserstein Anchor):

   • Antes de la fine-tuning, se calculan estadísticas gaussianas (media $\mu$ y varianza $\sigma^2$) de las activaciones del VLM pre-entrenado en una tarea de VQA (usando máscaras completas de imagen y texto).
   • Estas estadísticas definen un punto de referencia en el manifold de distribución de activaciones (manifold de Wasserstein-2).

2. Penalización de Transporte Wasserstein-2 ($L_{OT}$):

   • En cada paso de entrenamiento, se calcula la distancia $W_2^2$ entre las estadísticas actuales de activación y el ancla estática.
   • Esta distancia se descompone mediante la métrica de Bures en dos términos: desplazamiento de la media (mean shift) y desajuste de la desviación estándar (std mismatch).
   • Esta penalización genera un gradiente de referencia ($g_{ot}$) que indica la dirección necesaria para restaurar las activaciones al estado pre-entrenado.

3. Proyección Ortogonal de Gradientes por Capa (Layer-wise Orthogonal Gradient Projection):

   • Se realiza un doble retroceso secuencial: primero se extrae el gradiente de la tarea ($g_{task}$) y luego el gradiente del ancla ($g_{ot}$), manteniendo el grafo de cálculo.
   • Para cada capa del transformador, se aplica una proyección de Gram-Schmidt:
     • Si el gradiente de la tarea tiene una componente destructiva (producto punto negativo con el gradiente del ancla), se proyecta ortogonalmente.
     • Fórmula: $g_{final} = g_{task} - \alpha \cdot g_{ot}$, donde $\alpha$ es el coeficiente de proyección.
   • Granularidad: La proyección se realiza a nivel de capa (no global ni por tensor individual) para preservar la sincronización interna de los cabezados de atención (q, k, v, o) mientras se resuelve la interferencia en cada profundidad de la red.

3. Contribuciones Clave

1. Identificación de la Causa Raíz: Demostraron que la asimetría espectral (mismatch entre gradientes de alta dimensión semántica y baja dimensión física) es la causa geométrica del olvido, no solo un problema de capacidad de parámetros.
2. Marco sin Buffer (Buffer-Free): AEGIS elimina la necesidad de buffers de replay o co-entrenamiento con datos de VQA, utilizando únicamente un ancla estadística pre-calculada.
3. Intervención Geométrica Pura: A diferencia de métodos que restringen magnitudes (LoRA) o cortan gradientes (Stop-Gradient), AEGIS modifica solo la dirección del gradiente, eliminando componentes destructivos mientras preserva el 99%+ de la energía constructiva del gradiente.

4. Resultados Experimentales

Los experimentos se realizaron sobre PaliGemma2-3B en el benchmark LIBERO (tareas de manipulación en mesa), comparando AEGIS con Naive Fine-Tuning, Stop-Gradient + FAST, y LoRA.

• Preservación de VQA:
  • Naive: Pérdida de VQA severa (aumento de ~0.38 en la pérdida de holdout).
  • LoRA: Erosión lenta pero constante; no previene el olvido.
  • AEGIS: Preservación casi completa. La pérdida de VQA se mantiene cerca de la línea base (0.374 vs 0.392 inicial), superando a todos los métodos que permiten aprendizaje continuo.
• Aprendizaje de la Tarea (Acción):
  • AEGIS converge en la tarea de Flow-Matching tan rápido como la fine-tuning ingenua, demostrando que la proyección ortogonal no penaliza el aprendizaje de la nueva habilidad motora.
• Evaluación OOD (OK-VQA):
  • AEGIS mantuvo una precisión del 60.23% (casi idéntico al pre-entrenado 60.15%), mientras que la fine-tuning ingenua cayó a 57.36%.
• Eficiencia de la Proyección:
  • El sistema "desecha" (sheds) menos del 1% de la energía del gradiente en promedio (0.62%), confirmando que la interferencia destructiva es geométricamente delgada pero acumulativamente devastadora.
  • El costo computacional es un ~40% de sobrecarga en tiempo de pared (debido al segundo paso hacia atrás), mucho menor que el 100% de costo adicional del co-entrenamiento.

5. Significado e Impacto

El trabajo de AEGIS es significativo porque:

• Resuelve el dilema de la fine-tuning continua: Permite que los robots aprendan directamente de demostraciones de acciones continuas (sin tokenización discreta) sin sacrificar la inteligencia visual y lingüística pre-entrenada.
• Cambia el paradigma de protección: Mueve el enfoque de "restringir parámetros" (LoRA) o "replay de datos" a "control geométrico de la dirección del gradiente".
• Viabilidad Industrial: Al eliminar la necesidad de co-entrenamiento con datos de VQA y buffers de replay, reduce la complejidad logística y el costo computacional de entrenar modelos VLA a gran escala.

En conclusión, AEGIS demuestra que es posible alinear la geometría de los gradientes de control robótico con el manifold semántico pre-entrenado mediante proyecciones ortogonales inteligentes, logrando un equilibrio perfecto entre aprendizaje de nuevas habilidades y preservación de conocimiento.