Resumen Técnico: Generación Visual Autorregresiva de Árboles de Expansión (STAR)

Planteamiento del Problema

Los modelos autorregresivos (AR) basados en Transformers de solo Decodificador han alcanzado un éxito notable en la generación de lenguaje y se están aplicando cada vez más a la generación visual mediante la predicción del siguiente parche (next-patch prediction). Sin embargo, los modelos visuales AR convencionales suelen depender de un orden de escaneo rasterizado fijo (dependencia unidireccional). Aunque este orden ofrece un alto rendimiento de muestreo, impone una limitación estructural: el modelo no puede condicionarse a tokens que no aparecen como prefijos en la secuencia predefinida. Esta rigidez dificulta la edición nativa de imágenes (por ejemplo, inpainting), donde regiones arbitrarias de una imagen deben ser enmascaradas y completadas basándose en el contexto restante.

Para abordar esto, enfoques recientes emplean la permutación aleatoria de secuencias de tokens durante el entrenamiento para exponer a los modelos a diversos órdenes. Sin embargo, el artículo señala que entrenar con órdenes de secuencia arbitrarios suele derivar en un rendimiento de muestreo degradado en comparación con los modelos de escaneo rasterizado fijo. Esto crea un compromiso (trade-off): los órdenes fijos ofrecen rendimiento pero carecen de flexibilidad, mientras que las permutaciones aleatorias ofrecen flexibilidad pero sacrifican la calidad de generación y la velocidad de convergencia.

Metodología: Árbol de Expansión Autorregresivo (STAR)

Los autores proponen STAR, una estrategia de aleatorización de orden de secuencia estructurada que cierra esta brecha utilizando los órdenes de recorrido de Árboles de Expansión Uniformes (UST) sobre la red definida por las posiciones de los parches de la imagen.

1. Aleatorización Estructural mediante Árboles de Expansión Uniformes

En lugar de muestrear del espacio de todas las $N!$ permutaciones, STAR muestrea del espacio de árboles de expansión en una red de taxicab regular $G=(V, E)$, donde los vértices representan las posiciones de los parches.

• Selección de la Raíz: Se selecciona una esquina aleatoria de la imagen como nodo raíz $r$. Esta elección está motivada por la observación de que los órdenes de escaneo rasterizado (que comienzan desde las esquinas) suelen superar a otros, probablemente debido al sesgo de centro (los objetos salientes tienen menos probabilidades de aparecer en las esquinas) y la localidad (los parches adyacentes son más fáciles de predecir).
• Orden de Recorrido: La secuencia de orden $\tau$ se determina mediante un recorrido de Búsqueda en Anchura (BFS) del árbol de expansión muestreado.
• Base Teórica: Los autores hipotetizan que la ventaja de las secuencias predefinidas proviene de la preservación del conocimiento previo (sesgo de centro y localidad). Verifican esto midiendo la entropía del modelo condicional, demostrando que los tokens más cercanos a los tokens predichos previamente (en términos de distancia del árbol) exhiben una menor entropía, de forma similar al comportamiento del escaneo rasterizado. El número de árboles de expansión crece exponencialmente ($\sim \exp(N \cdot z_G)$) pero sigue siendo un subconjunto mucho más pequeño y estructurado del espacio completo de permutaciones.

2. Inferencia y Completado de Postfijo (Inpainting de Imágenes)

Un desafío crítico para los modelos AR en inpainting es asegurar que una observación parcial dada (la región no enmascarada) aparezca como un prefijo en el orden de la secuencia, permitiendo que la región enmascarada se genere como un postfijo.

• Muestreo por Rechazo: A diferencia de la permutación aleatoria, donde la separación prefijo/postfijo es trivial, STAR debe muestrear un árbol de expansión donde el recorrido BFS coloque naturalmente la región no enmascarada antes que la región enmascarada.
• Algoritmo:
  1. Identificar los límites de la región no enmascarada.
  2. Seleccionar un nodo raíz $r$ que sea el más lejano (en distancia de Manhattan) del límite más cercano de la región no enmascarada para maximizar la probabilidad de que la región no enmascarada sea visitada primero.
  3. Muestrear un árbol de expansión en la región no enmascarada. Si la profundidad máxima del árbol no satisface la condición de que todos los vértices del límite sean visitados antes que cualquier vértice enmascarado, la muestra se rechaza (muestreo por rechazo).
  4. Una vez encontrado un árbol válido para la región no enmascarada, se muestrean árboles de expansión para las regiones enmascaradas y se conectan a los vértices del límite para completar el árbol global.
• Eficiencia: Los autores demuestran que el uso de BFS con la selección de "raíz más lejana" reduce significamente el número de ensayos de rechazo en comparación con la Búsqueda en Profundidad (DFS), haciendo que la inferencia sea eficiente.

3. Arquitectura del Modelo

STAR requiere cambios mínimos en la arquitectura estándar de Transformer de solo Decodificador. Incorpora:

• Máscaras de atención causal.
• Características adicionales para condicionarse en la posición del siguiente token (usando embeddings posicionales aprendibles), de forma similar a métodos previos basados en permutaciones.
• No se requieren cambios arquitectónicos complejos (por ejemplo, atención bidireccional o tokenizadores especializados).

Contribuciones Clave

1. Estrategia de Aleatorización Estructurada: La propuesta del uso del recorrido BFS de árboles de expansión uniformes como orden de secuencia para la generación visual AR. Esta estrategia preserva estructuralmente los priors de la imagen (sesgo de centro, localidad) mientras introduce suficiente aleatoriedad para la flexibilidad.
2. Compromiso entre Rendimiento y Flexibilidad: La demostración de que el recorrido BFS de los UST es suficiente para lograr los beneficios de la permutación aleatoria (diversos órdenes de secuencia) sin la degradación de rendimiento típicamente asociada con las permutaciones arbitrarias. Esto se logra sin aumentar la complejidad del entrenamiento.
3. Soporte Nativo para Edición de Imágenes: La introducción de un mecanismo (muestreo por rechazo con BFS) que permite la modificación a nivel de región (inpainting) con cambios arquitectónicos mínimos, abordando una limitación clave de los modelos AR convencionales.

Resultados Experimentales

Los autores evalúan STAR en ImageNet-1k (resolución 256x256) utilizando tareas de generación condicionada por clase e inpainting de imágenes.

• Generación de Imágenes: Los modelos STAR (STAR-B, L, XL, XXL) logran un rendimiento comparable a modelos de vanguardia (por ejemplo, RAR, DiT, LlamaGen) en términos de Distancia de Incepción de Fréchet (FID) y Inception Score (IS). Notablemente, STAR-XXL logra un FID de 1.55, superando a RAR-XXL (1.48) y RandAR-XXL (2.15), manteniendo puntuaciones competitivas.
• Inpainting de Imágenes: STAR supera significativamente tanto a los modelos de orden fijo (RAR) como a los modelos de permutación aleatoria (RandAR) en tareas de inpainting.
  • En varios ratios de enmascaramiento (0.1 a 0.9), STAR mantiene altos valores de FID e IS.
  • En contraste, RAR (escaneo rasterizado fijo) muestra una caída significativa en el rendimiento de inpainting porque no puede manejar naturalmente máscaras arbitrarias como prefijos.
  • RandAR mantiene la capacidad de inpainting pero sufre de una menor calidad de generación en comparación con STAR.
• Estudios de Ablación:
  • Localidad: La permutación aleatoria no logra superar a STAR, lo que sugiere que los órdenes de secuencia que violan la localidad son perjudiciales para la convergencia.
  • Sesgo de Centro: La selección aleatoria de esquinas como raíces funciona bien para imágenes con un fuerte sesgo de centro y generaliza mejor que la permutación aleatoria para imágenes con un sesgo de centro más débil.
  • Análisis de Entropía: Los modelos entrenados con STAR exhiben una escala más uniforme de la entropía del modelo condicional a través de las posiciones de los tokens en comparación con la permutación aleatoria, lo cual se correlaciona con una mejor convergencia.
  • Sobrecarga: La sobrecarga computacional para el muestreo de árboles de expansión es insignificante (< 0.0004% del tiempo total de inferencia). El muestreo por rechazo para inpainting es altamente eficiente, requiriendo típicamente menos de 5 ensayos con BFS.

Significación y Reivindicaciones

El artículo afirma que STAR ofrece una solución simple pero efectiva al compromiso de larga duración entre el rendimiento de muestreo y la flexibilidad del orden de secuencia en la generación visual AR. Al aprovechar las propiedades estructurales de los árboles de expansión uniformes, STAR:

• Preserva la arquitectura Transformer de solo Decodificador, facilitando la integración en modelos multimodales de fusión temprana.
• Evita la necesidad de estrategias de entrenamiento complejas como el anidamiento de orden (order annealing, cambiando de órdenes aleatorios a fijos durante el entrenamiento) o tokenizadores especializados.
• Proporciona soporte nativo para la edición de imágenes (inpainting) sin sacrificar la calidad de generación.

Los autores posicionan a STAR como un "bloque de construcción básico" para futuros modelos visuales AR, enfatizando que su simplicidad permite que sea fácilmente adoptado y extendido a modelos generativos multimodales. El trabajo sugiere que el éxito de los modelos AR depende no solo de la arquitectura del modelo, sino del orden de secuencia de aleatorización estructurada para alinearse con las propiedades intrínsecas de los datos visuales.