Where Bits Matter in World Model Planning: A Paired Mixed-Bit Study for Efficient Spatial Reasoning

1. Problema y Contexto

El razonamiento espacial eficiente mediante modelos de mundo (world models) es crucial para el control y la planificación, especialmente cuando se basan en características visuales preentrenadas (como DINO-WM). Sin embargo, las restricciones de despliegue (memoria y latencia) obligan a utilizar cuantización agresiva de bajo ancho de bits.

El problema central que abordan los autores es determinar si la degradación del rendimiento en entornos de planificación de bajo ancho de bits se debe principalmente a:

1. El ancho de bits total promedio del modelo.
2. La distribución de bits entre los diferentes módulos del modelo (específicamente, entre el codificador visual y el predictor de dinámicas latentes).

La hipótesis es que, cerca de los puntos de operación de bajo ancho de bits, la calidad de la planificación no depende solo de la precisión media, sino de dónde se asignan los bits.

2. Metodología

Los autores realizaron un estudio empírico utilizando el modelo DINO-WM en la tarea de planificación "Wall" (con obstáculos y objetivos espaciales).

• Configuración Experimental:
  • Se evaluaron variantes uniformes (todos los módulos con el mismo ancho de bits) y variantes mixtas/asimétricas (diferentes anchos de bits para el codificador y el predictor).
  • Variantes probadas: FP16 (precisión completa), INT8, INT6, INT4 e INT3 (uniformes). Variantes mixtas donde el codificador se mantiene en FP16 o INT8/6/4 y el predictor se cuantiza, y viceversa.
  • Protocolo de Evaluación: Se utilizó una evaluación de objetivos emparejados (paired-goal evaluation) para controlar la variabilidad. Se probaron dos presupuestos de planificador:
    • Presupuesto bA: 9 pasos de horizonte, 2 pasos óptimos, 2 iteraciones máximas (3 semillas).
    • Presupuesto bB: 12 pasos de horizonte, 3 pasos óptimos, 3 iteraciones máximas (2 semillas).
  • Implementación: Cuantización post-entrenamiento (PTQ) basada solo en pesos (weight-only) en capas lineales, sin recalibración ni reentrenamiento.

3. Contribuciones Clave

El estudio identifica tres regímenes de comportamiento distintos en la relación entre el ancho de bits y el éxito de la planificación:

1. Régimen Estable (8/6 bits): Los modelos con 8 o 6 bits (tanto uniformes como mixtos) mantienen un rendimiento cercano al de FP16.
2. Régimen de Colapso (3 bits): Todas las variantes de 3 bits colapsan completamente, logrando un 0% de éxito en la planificación.
3. Régimen de Transición Sensible (4 bits): Este es el hallazgo más crítico. En 4 bits, el rendimiento es altamente sensible a la asignación de bits.
   • Las variantes mixtas (codificador de alta precisión, predictor de 4 bits) superan consistentemente a las variantes uniformes de 4 bits.
   • Esto sugiere que preservar la precisión del codificador visual es más crítico que la del predictor en este umbral.

4. Resultados Principales

• Comparación Uniforme vs. Mixta (INT4):

  • En el presupuesto bA, la variante mixta INT4 logró un 26.7% de éxito frente al 6.7% de la uniforme INT4.
  • En el presupuesto bB, la variante mixta INT4 alcanzó un 50.0% de éxito frente al 20.0% de la uniforme.
  • Aunque la diferencia no fue estadísticamente definitiva en todas las pruebas (debido al tamaño de muestra), la dirección del efecto fue consistente.

• Análisis de Asimetría y Tamaño:

  • Para descartar que la mejora se debiera simplemente a un mayor tamaño total del modelo (ya que las variantes mixtas suelen ser más grandes), se compararon variantes asimétricas de tamaño similar a la uniforme INT4.
  • La variante E6/P4 (Codificador INT6, Predictor INT4) superó a la uniforme INT4, mientras que la variante E4/P8 (Codificador INT4, Predictor INT8) no logró recuperar el rendimiento.
  • Conclusión: Invertir bits en el codificador es más efectivo que invertirlos en el predictor cerca del límite de 4 bits.

• Validación estricta (Replicación):

  • Una réplica estricta con menos episodios confirmó la estructura de regímenes (8 bits > 4 bits > 3 bits), pero mostró que la ventaja de la variante mixta sobre la uniforme en 4 bits puede ser condicional al presupuesto de planificación en muestras pequeñas, destacando la sensibilidad de este régimen de transición.

• Diagnóstico Mecanístico:

  • Se encontró una fuerte correlación negativa (ρ = -0.928) entre la divergencia de las representaciones latentes y el éxito de la planificación. Esto sugiere que la degradación del codificador en bajo ancho de bits daña la geometría latente necesaria para el razonamiento espacial.

5. Significado y Direcciones Futuras

• Cambio de Paradigma: El trabajo sugiere que la asignación de bits no debe tratarse como un detalle de compresión post-hoc, sino como un problema de asignación de recursos consciente de la planificación.
• Políticas de Cuantización: Para la planificación espacial eficiente, las políticas de cuantización deben ser "conscientes del módulo" (module-aware). Priorizar la precisión del codificador visual es esencial cuando el presupuesto de bits es limitado (zona de transición de 4 bits).
• Investigación Futura: Los autores proponen optimizar conjuntamente la precisión del codificador y el predictor bajo restricciones explícitas de memoria y latencia. También sugieren explorar la cuantización consciente de activaciones, la calibración y la localización a nivel de bloques o cabezas de atención para refinar estas políticas.

En resumen, el paper demuestra que en los modelos de mundo para planificación, la ubicación de los bits importa tanto como el número total de bits, y que preservar la integridad del codificador visual es la estrategia clave para mantener la capacidad de razonamiento espacial en configuraciones de bajo ancho de bits.