Structured State-Space Regularization for Compact and Generation-Friendly Image Tokenization

Resumen Técnico: Regularización de Espacio de Estados Estructurado para Tokenización de Imágenes Compacta y Amigable para la Generación

1. El Problema

Los modelos de visión modernos (generativos, de mundo y de representación) dependen cada vez más de tokenizadores de imágenes que operan en espacios latentes. Sin embargo, existe un dilema fundamental en el diseño de estos espacios latentes:

• Compacidad vs. Preservación de Información: El espacio latente debe ser compacto (baja dimensionalidad) para reducir costos computacionales, pero debe retener suficiente información esencial para una reconstrucción fiel.
• Amigabilidad para la Generación (Generation-Friendliness): El espacio latente debe ser fácil de modelar para arquitecturas generativas modernas (como los Modelos de Difusión).

Actualmente, existe una compensación (trade-off): los tokenizadores que logran una reconstrucción perfecta a menudo generan espacios latentes difíciles de modelar para la difusión, mientras que aquellos optimizados para la generación pueden perder fidelidad en la reconstrucción. El objetivo es encontrar un espacio latente que sea compacto, informativo y estructurado de manera que favorezca la síntesis progresiva de frecuencias (de baja a alta), una propiedad intrínseca de los modelos de difusión.

2. Metodología Propuesta

Los autores proponen una nueva técnica de regularización llamada Regularización de Espacio de Estados Estructurado (Structured State-Space Regularization). La idea central es alinear la dinámica del espacio latente del tokenizador con la dinámica de los Modelos de Espacio de Estados (SSM), específicamente transfiriendo su propiedad de "conciencia de frecuencia".

Conceptos Clave de la Metodología:

• Reinterpretación de los SSM: Los autores argumentan que los SSM pueden verse como una combinación de dos componentes:

  1. Proyección de Base ($c(\cdot)$): Proyectar la señal de entrada sobre una base de funciones (ej. polinomios ortogonales o bases de Fourier).
  2. Transformación de Entrada ($\theta(\cdot)$): Una transformación aplicada a la entrada a lo largo del tiempo.

• Selección de Transformación: En lugar de la concatenación de tokens típica de los SSM 1D, el método utiliza un proceso de desenfoque gaussiano progresivo ($\theta$) sobre la imagen. Esto es teóricamente sólido porque el desenfoque gaussiano es la solución a la ecuación del calor, lo que permite una descomposición analítica de las frecuencias (eliminando frecuencias altas progresivamente).

• Derivación de la Dinámica:

  • Se define una proyección de la imagen $I$ sobre una base 2D ortogonal (ej. Fourier) para obtener coeficientes $c(I)$.
  • Se demuestra matemáticamente que la evolución de estos coeficientes bajo un desenfoque gaussiano sigue una dinámica lineal gobernada por una matriz de transición de estado $\mathbf{A}$:
    $$\frac{d}{d\tau} c(I_\tau) = \mathbf{A} c(I_\tau)$$
  • Al discretizar esta ecuación, se obtiene una regla de actualización: $c(I_t) = \bar{\mathbf{A}} c(I_{t-1})$.

• Función de Pérdida de Regularización:
  Se entrena un tokenizador (codificador $E$ y decodificador $D$) para que sus representaciones latentes imiten esta dinámica. Dadas dos imágenes con diferentes niveles de desenfoque ($I_{\tau_1}$ y $I_{\tau_2}$):

  1. Se codifica la imagen menos borrosa: $z_1 = E(I_{\tau_1})$.
  2. Se aplica la matriz de transición aprendida/definida: $\hat{z}_2 = \bar{\mathbf{A}} z_1$.
  3. Se codifica la imagen más borrosa: $z_2 = E(I_{\tau_2})$.
  4. Se minimiza la distancia entre la predicción $\hat{z}_2$ y el objetivo real $z_2$, además de incluir una pérdida de reconstrucción del decodificador para evitar colapso de modos.

  $$\mathcal{L} = \lambda_z d_z(z_2, \bar{\mathbf{A}}z_1) + \lambda_I d_I(I_{\tau_2}, D(\bar{\mathbf{A}}z_1))$$

• Centrado Espacial: Para evitar que las normas de los vectores latentes decaigan a cero (debido a la naturaleza del desenfoque), se aplica una función de centrado espacial $\mathcal{C}$ antes de la regularización, estabilizando el entrenamiento.

3. Contribuciones Clave

1. Nuevo Regularizador Teórico: Propone un regularizador motivado teóricamente que guía a los tokenizadores a adoptar la estructura de frecuencia inherente de los SSM.
2. Marco Teórico SSM-Imagen: Presenta el primer marco que integra principios de SSM (específicamente la proyección de bases y la dinámica de coeficientes bajo transformaciones de imagen) en la tokenización de imágenes 2D.
3. Eficacia Comprobada: Demuestra que este enfoque mejora significativamente la capacidad de generación de los modelos de difusión sin sacrificar significativamente la fidelidad de reconstrucción.

4. Resultados Experimentales

Los autores evaluaron su método en dos tokenizadores populares: Flux (ratio de submuestreo 8x) y Cosmos (ratio 16x), utilizando el conjunto de datos ImageNet-1K.

• Rendimiento de Generación:
  • En el tokenizador Flux, el método propuesto redujo el gFID (Fréchet Inception Distance de generación) de 16.08 (base) a 13.33, superando a métodos anteriores como EQVAE (13.95) y FT-SE (13.87).
  • El IS (Inception Score) mejoró drásticamente de 89.32 a 106.61.
  • En el tokenizador Cosmos, se logró un gFID de 11.20 (frente a 13.52 del base), con mejoras consistentes en todas las métricas de generación.
• Rendimiento de Reconstrucción:
  • La pérdida en la calidad de reconstrucción fue mínima. Por ejemplo, en Flux, el PSNR bajó ligeramente de 31.70 a 31.55, manteniendo métricas como SSIM y LPIPS en niveles comparables o ligeramente mejores que las baselines.
• Análisis de Estructura Latente:
  • Las visualizaciones mostraron que los canales latentes aprendieron a codificar bandas de frecuencia específicas. Al revelar los canales de baja a alta frecuencia, la imagen reconstruida se formó progresivamente (de contornos generales a detalles finos), confirmando que el espacio latente ha adquirido una estructura espectral ordenada.
  • La dinámica de los canales latentes imitó la función exponencial de los coeficientes de Fourier bajo desenfoque, validando la hipótesis teórica.

5. Significado e Impacto

Este trabajo es significativo por varias razones:

• Puente entre Dominios: Conecta dos campos aparentemente desconectados: los Modelos de Espacio de Estados (SSM), conocidos por su eficiencia en secuencias y conciencia de frecuencia, y los Tokenizadores de Imágenes para modelos generativos.
• Resolución del Dilema Reconstrucción-Generación: Ofrece una solución elegante al problema de compensación entre reconstrucción fiel y facilidad de generación, demostrando que imponer una estructura espectral correcta en el espacio latente beneficia a ambas tareas.
• Nueva Perspectiva de Diseño: Sugiere que la "amigabilidad para la generación" no es solo una propiedad emergente, sino una característica que puede ser inducida mediante regularización basada en principios físicos y matemáticos (como la ecuación del calor y bases ortogonales).
• Aplicabilidad: El método es compatible con arquitecturas de tokenizadores existentes (como VQ-VAE o Autoencoders Variacionales) y se puede implementar con un costo computacional adicional bajo durante el entrenamiento.

En conclusión, la regularización de espacio de estados estructurada proporciona una vía prometedora para crear espacios latentes que no solo comprimen imágenes eficientemente, sino que también organizan la información de manera que los modelos generativos modernos puedan sintetizar imágenes de alta calidad de forma más robusta.