💻 computer science

MOSAIK: Multi-Patch Content-Aware Spatial Allocation of Image Tokens for Efficient Generation

MOSAIK es un marco de trabajo guiado por daño para modelos de difusión en el espacio de píxeles que asigna dinámicamente tamaños de parches heterogéneos a través de regiones de la imagen basándose en la pérdida de fidelidad estimada, logrando una eficiencia computacional significativa (reducción del 70% de FLOPs) mientras mantiene una calidad de generación competitiva en comparación con las líneas base de cómputo completo.

Autores originales: Mohammadreza Hami, Mohammadreza Samadi, Chao Gao, Negar Hassanpour

Publicado 2026-08-07

📖 4 min de lectura☕ Lectura para el café

CC BY 4.0

Autores originales: Mohammadreza Hami, Mohammadreza Samadi, Chao Gao, Negar Hassanpour

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás intentando pintar una obra maestra en un lienzo gigante, pero solo tienes una cantidad limitada de pintura y un límite de tiempo muy estricto. En el mundo de la inteligencia artificial, crear imágenes es un poco como esto. Durante mucho tiempo, las computadoras usaron un método "comprimido" para dibujar imágenes, algo así como bosquejar una escena en un cuaderno diminuto antes de intentar rellenar los detalles. Esto era rápido, pero tenía un techo: una vez que la habías bosquejado pequeña, no podías recuperar mágicamente los detalles diminutos y nítidos, como la textura del pelaje o los pétalos individuales de una flor.

Para superar esto, una nueva generación de IA comenzó a pintar directamente sobre el lienzo gigante, píxel por píxel. Esto es como cambiar de un cuaderno diminuto a un mural masivo. ¿El resultado? Imágenes asombrosamente realistas. Pero hay un inconveniente: pintar cada uno de los píxeles requiere que la computadora procese una cantidad masiva de información, lo cual es lento y costoso. Para acelerar esto, los científicos han intentado usar "pinceles grandes" para cubrir áreas extensas del lienzo con un solo trazo. Sin embargo, usar un pincel grande en todas partes es un poco como intentar pintar el retrato de un tigre con un rodillo de pintor de casas; podrías lograr la forma general, pero perderías las rayas y los bigotes. La gran pregunta ha sido: ¿Cómo usamos pinceles grandes para el fondo aburrido (como el cielo) pero pinceles diminutos y precisos para las partes importantes (como la cara del tigre) sin que la computadora se confunda?

Aquí es donde entra un nuevo método llamado MOSAIK. Piensa en MOSAIK como un director de arte superinteligente para el pintor de IA. En lugar de obligar a la computadora a usar el mismo tamaño de pincel para toda la imagen, MOSAIK observa la imagen mientras se está creando y decide: "Oye, esta parte es solo un cielo borroso, usemos un pincel gigante aquí para ahorrar tiempo. Pero esta parte es el ojo del tigre, cambiemos a un pincel diminuto y súper preciso ahora mismo".

El artículo presenta MOSAIK (que significa Asignación Espacial de Tokens de Imagen de Contenido Multiparche) como una forma de hacer que estos generadores de imágenes de alta calidad, píxel por píxel, sean mucho más rápidos sin arruinar la imagen. Los investigadores descubrieron que, al cambiar dinámicamente el tamaño de las "pinceladas" (o parches) a través de diferentes partes de la imagen, pudieron reducir la carga de trabajo de la computadora en una cantidad enorme. Específicamente, lograron reducir el número de cálculos (llamados FLOPs) en un 70% y el número de fragmentos de datos (tokens) en un 83%.

Aquí está el truco de magia: MOSAIK no solo adivina dónde usar pinceles grandes. Utiliza un "predictor de daño". Imagina que la IA tiene un sensor especial que pregunta: "Si uso un pincel grande en este punto específico, ¿cuánta información perderé?". Si la respuesta es "mucha" (como en el pelaje de un tigre), mantiene el pincel pequeño. Si la respuesta es "nada" (como en un cielo azul suave), cambia a un pincel grande. El artículo muestra que este enfoque inteligente e irregular funciona mucho mejor que los métodos anteriores que intentaban acelerar las cosas simplemente usando pinceles grandes para toda la imagen en ciertos momentos o saltándose pasos.

En sus pruebas, MOSAIC fue capaz de generar imágenes que se veían casi exactamente iguales a la versión lenta de máxima potencia, a pesar de estar haciendo un 70% menos de trabajo. Cuando compararon MOSAIK con otros trucos de aceleración, MOSAIK mantuvo los detalles nítidos mientras que los otros empezaban a verse borrosos o perdían características importantes. Los investigadores sugieren que este enfoque "guiado por el daño" es una forma poderosa de hacer que la generación de imágenes por IA de alta calidad sea más rápida y eficiente, demostrando que no tienes que elegir entre velocidad y calidad si sabes dónde ser eficiente y dónde ser preciso.

Resumen Técnico: MOSAIK

Planteamiento del Problema

Los modelos de difusión en el espacio de píxeles (PDMs) ofrecen una ventaja significativa sobre los modelos de difusión latente (LDMs) al evitar el techo de reconstrucción impuesto por los Autoencoders Variacionales (VAEs), preservando así detalles espaciales de alta frecuencia y grano fino. Sin embargo, generar directamente en el espacio de píxeles resulta en un recuento de tokens sustancialmente mayor. Dado que los mecanismos de autoatención en los transformers exhiben una complejidad cuadrática con respecto al recuento de tokens, los PDMs incurren en costos computacionales prohibitivos.

Los métodos de eficiencia existentes intentan mitigar esto variando los tamaños de parche (el área espacial representada por un único token) a través de los pasos de eliminación de ruido o la profundidad de la red. Aunque estos enfoques reducen el recuento total de tokens, aplican un diseño de parches espacialmente uniforme en toda la imagen en cualquier paso dado. Esta uniformidad crea una ineficiencia fundamental: o bien gasta computación en exceso en regiones simples y de baja información, o bien restringe la capacidad de representación en regiones sensibles donde los detalles finos son críticos. El artículo argumenta que la densidad de la información visual es inherentemente heterogénea, lo que requiere una estrategia de asignación dinámica por región que los métodos actuales carecen.

Metodología: MOSAIK

MOSAIK (Multi-patch cOntent-aware Spatial Allocation of Image toKens) es un marco diseñado para variar los tamaños de parche a través de diferentes regiones de la imagen y pasos de eliminación de ruido basados en el "daño" predicho (pérdida de fidelidad) causado por el engrosamiento. El enfoque se basa en la arquitectura base PixelDiT y utiliza un flujo de entrenamiento de tres etapas:

Etapa 1: Fluidez de Múltiples Parches

El objetivo es equipar al modelo para operar en múltiples tamaños de parche uniformes ( $P \in \{16, 32, 64\}$ ) sin un olvido catastrófico del prior generativo original.

Arquitectura: Los pesos de la arquitectura base están congelados. Se inyectan módulos de Adaptación de Bajo Rango (LoRA) en los bloques DiT. Se introducen tokenizadores separados para $p32$ y $p64$ , inicializados mediante la composición de la proyección $p16$ preentrenada con submuestreo bilineal.
Entrenamiento: El modelo se somete a una autodestilación, optimizando para igualar las predicciones de velocidad del modelo profesor $p16$ congelado a través de todos los tamaños de parche uniformes.

Etapa 2: Adaptación de Parches Mixtos

Esta etapa permite al modelo procesar diseños de parches arbitrarios y heterogéneos (donde diferentes regiones usan diferentes tamaños de parche) dentro de un solo pase hacia adelante (forward pass).

Embeddings Posicionales: Los Rotary Positional Embeddings (RoPE) estándar asumen una cuadrícula uniforme. MOSAIK aborda el desajuste espacial de los parches mixtos anclando cada token al centro geométrico de su parche dentro de un marco de coordenadas $p16$ normalizado.
Entrenamiento: El generador se entrena con diseños heterogéneos muestreados aleatoriamente, desacoplando deliberadamente la disposición espacial de los tamaños de parche del contenido de la imagen. Esto obliga al modelo a aprender a ejecutar cualquier diseño independientemente de su alineación con las características visuales.

Etapa 3: Predictor de Diseño Guiado por Daño

Esta etapa introduce un predictor ligero que asigna dinámicamente los tokens basándose en la degradación estimada causada por el engrosamiento de regiones específicas.

Definición de Daño: El "daño" se define como la diferencia al cuadrado entre la predicción de un diseño uniforme $p64$ y el profesor de alta fidelidad $p16$ , promediada sobre regiones no superpuestas de $64 \times 64$ .
Entrenamiento del Predictor: Una red convolucional ligera con condicionamiento FiLM se entrena para regredir este mapa de daño directamente desde las activaciones intermedias de la arquitectura base (específicamente las activaciones del bloque-6 de DiT) y los embeddings de texto. Esto evita la necesidad de un segundo pase hacia adelante durante la inferencia.
Estrategia de Inferencia: En cada paso de eliminación de ruido, el predictor estima el mapa de daño. Las regiones con alto daño predicho se asignan parches más finos ( $p16$ ), mientras que las regiones de bajo daño se asignan parches más gruesos ( $p32$ o $p64$ ). La asignación está restringida por un presupuesto de tokens especificado por el usuario, priorizando los refinamientos que ofrecen la mayor reducción de daño por cada token añadido.

Contribuciones Clave

Parcheificación Espacialmente Adaptativa: MOSAIK es el primer marco que aplica diseños de parches heterogéneos dentro de un modelo de difusión en el espacio de píxeles, yendo más allá de las restricciones de uniformidad espacial de los métodos de eficiencia previos.
Asignación Guiada por Daño: Los autores formulan la asignación de parches como un problema de predicción de daño, utilizando un predictor ligero para identificar las regiones donde el engrosamiento causaría la pérdida de fidelidad más significativa.
Flujo de Entrenamiento de Tres Etapas: Una estrategia de entrenamiento novedosa que primero establece la fluidez de múltiples parches, luego se adapta a diseños mixtos y, finalmente, entrena un predictor separado para guiar la ruta de los tokens sin reentrenar el generador.
Capacidad de Integración: El marco está diseñado para ser ortogonal a los métodos de eficiencia temporal (como el almacenamiento en caché de características), permitiendo una optimización combinada.

Resultados

Los experimentos se realizaron utilizando generación de imágenes de 1024x1024 en los entornos de prueba GenEval y DPG-Bench, comparando MOSAIK contra líneas base que incluyen reducción de parche uniforme, programación de parches temporal (DDiT) y almacenamiento en caché de características (TeaCache, TaylorSeer).

Eficiencia vs. Calidad: MOSAIK logra una reducción del 70% en FLOPs y una reducción del 83% en el recuento de tokens mientras iguala al PixelDiT de cómputo completo en GenEval (puntuación 0.74) y solo cae 1.0 punto en DPG-Bench (82.5 vs. 83.5).
Comparación con la Programación Temporal: En regímenes de cómputo agresivos (por ejemplo, reducción del 60-70% de FLOPs), MOSAIK supera consistentemente a DDiT. Mientras que DDiT aplica un único tamaño de parche globalmente, MOSAIK redirige dinámicamente el cómputo hacia las regiones sensibles (por ejemplo, límites de sujetos), preservando la integridad estructural y los detalles finos como texturas y bordes.
Comparación con el Almacenamiento en Caché: MOSAIK supera a las líneas base basadas en caché (TeaCache, TaylorSeer) en presupuestos equivalentes, particularmente en regímenes altamente restringidos.
Complementariedad: Al combinarse con métodos de caché, MOSAIK mantiene el rendimiento en niveles de compresión aún más altos (hasta un 80% de reducción de FLOPs), donde el caché por sí solo causa un colapso significativo de la calidad.
Ablación: Reemplazar el predictor guiado por daño con un asignador aleatorio o ciego al contenido resulta en una degradación significativa del rendimiento, validando que el ranking de daño predicho conlleva información genuina sobre la sensibilidad espacial.

Significado y Reivindicaciones

El artículo afirma que MOSAIK demuestra que la heterogeneidad espacial en la información visual es un factor crítico para la generación eficiente que ha sido pasado por alto por los métodos centrados únicamente en la adaptación temporal o de profundidad.

Asignación Óptima de Recursos: Los autores argumentan que la parcheificación fina uniforme desperdicia computación en regiones simples, mientras que la parcheificación gruesa uniforme destruye detalles en regiones complejas. MOSAIK cierra esta brecha asignando recursos donde son más necesarios.
Ventaja del Espacio de Píxeles: El marco aprovecha la redundancia espacial inherente disponible en los modelos de espacio de píxeles (a diferencia de los modelos latentes donde la compresión del VAE ya ha consumido la redundancia) para utilizar dinámicamente el engrosamiento sin la necesidad de una optimización conjunta de extremo a extremo de un enrutador y un generador.
Escalabilidad: El enfoque permite un rendimiento estable a través de una amplia gama de presupuestos de cómputo, desde reducciones moderadas hasta restricciones agresivas, sin requerir el reentrenamiento para diferentes niveles de presupuesto.

El artículo concluye que MOSAIK establece un nuevo estándar para la eficiencia en la difusión en el espacio de píxeles, ofreciendo una alternativa altamente competitiva a la programación temporal y al almacenamiento en caché, y sugiere que el trabajo futuro podría extender estos principios a la difusión de video, explotando la redundancia tanto en el espacio como en el tiempo.