Resumo Técnico: MOSAIK

Definição do Problema

Modelos de difusão no espaço de pixels (PDMs) oferecem uma vantagem significativa sobre os modelos de difusão latente (LDMs) ao evitar o teto de reconstrução imposto pelos Autoencoders Variacionais (VAEs), preservando, assim, detalhes espaciais de alta frequência e granulação fina. No entanto, gerar diretamente no espaço de pixels resulta em uma contagem de tokens substancialmente maior. Como os mecanismos de autoatenção em transformers exibem complexidade quadrática em relação à contagem de tokens, os PDMs incorrem em custos computacionais proibitivos.

Métodos de eficiência existentes tentam mitigar isso variando os tamanhos de patch (a área espacial representada por um único token) através dos passos de denoising ou da profundidade da rede. Embora essas abordagens reduzam a contagem total de tokens, elas aplicam um layout de patch espacialmente uniforme em toda a imagem em qualquer etapa. Essa uniformidade cria uma ineficiência fundamental: ou gasta-se computação excessiva em regiões simples e de baixa informação ou, inversamente, restringe-se a capacidade representacional em regiões sensíveis onde detalhes finos são críticos. O artigo argumenta que a densidade de informação visual é inerentemente heterogênea, necessitando de uma estratégia de alocação dinâmica por região que os métodos atuais não possuem.

Metodologia: MOSAIK

O MOSAIK (Multi-patch cOntent-aware Spatial Allocation of Image toKens) é um framework projetado para variar os tamanhos de patch entre diferentes regiões da imagem e passos de denoising com base no "dano" previsto (perda de fidelidade) causado pelo coarse-graining (redução de resolução). A abordagem é construída sobre o backbone PixelDiT e utiliza um pipeline de treinamento de três estágios:

Estágio 1: Fluência de Múltiplos Patches (Multi-Patch Fluency)

O objetivo é equipar o modelo para operar em múltiplos tamanhos de patch uniformes ($P \in \{16, 32, 64\}$) sem o esquecimento catastrófico do prior generativo original.

• Arquitetura: Os pesos do backbone são congelados. Módulos de Adaptação de Baixo Posto (LoRA) são injetados nos blocos DiT. Tokenizadores separados são introduzidos para $p32$ e $p64$, inicializados pela composição da projeção $p16$ pré-treinada com downsampling bilinear.
• Treinamento: O modelo passa por autodestilação, otimizando para corresponder às previsões de velocidade do modelo professor $p16$ congelado em todos os tamanhos de patch uniformes.

Estágio 2: Adaptação de Patches Mistos (Mixed-Patch Adaptation)

Este estágio permite que o modelo processe layouts de patches arbitrários e heterogêneos (onde diferentes regiões usam diferentes tamanhos de patch) dentro de um único forward pass.

• Embeddings Posicionais: Os Rotary Positional Embeddings (RoPE) padrão assumem uma grade uniforme. O MOSAIK aborda o descompasso espacial de patches mistos ancorando cada token ao centro geométrico de seu patch dentro de um quadro de coordenadas $p16$ normalizado.
• Treinamento: O gerador é treinado em layouts heterogêneos amostrados aleatoriamente, desacoplando deliberadamente o arranjo espacial dos tamanhos de patch do conteúdo da imagem. Isso força o modelo a aprender a executar qualquer layout, independentemente de seu alinhamento com características visuais.

Estágio 3: Preditor de Layout Guiado por Dano (Damage-Guided Layout Predictor)

Este estágio introduz um preditor leve que aloca dinamicamente tokens com base na degradação estimada causada pelo coarse-graining de regiões específicas.

• Definição de Dano: O "dano" é definido como a diferença ao quadrado entre a previsão de um layout uniforme $p64$ e o professor de alta fidelidade $p16$, calculada pela média sobre regiões não sobrepostas de $64 \times 64$.
• Treinamento do Preditor: Uma rede convolucional leve com condicionamento FiLM é treinada para regredir este mapa de dano diretamente das ativações intermediárias do backbone (especificamente as ativações do bloco DiT-6) e dos embeddings de texto. Isso evita a necessidade de um segundo forward pass durante a inferência.
• Estratégia de Inferência: Em cada passo de denoising, o preditor estima o mapa de dano. Regiões com alto dano previsto recebem patches mais finos ($p16$), enquanto regiões de baixo dano recebem patches mais grosseiros ($p32$ ou $p64$). A alocação é restrita por um orçamento de tokens especificado pelo usuário, priorizando refinamentos que ofereçam a maior redução de dano por token adicionado.

Principais Contribuições

1. Patchificação Espacialmente Adaptativa: O MOSAIK é o primeiro framework a aplicar layouts de patches heterogêneos dentro de um modelo de difusão no espaço de pixels, indo além das restrições de uniformidade espacial de métodos anteriores.
2. Alocação Guiada por Dano: Os autores formulam a alocação de patches como um problema de previsão de dano, usando um preditor leve para identificar regiões onde o coarse-graining causaria a perda de fidelidade mais significativa.
3. Pipeline de Treinamento de Três Estágios: Uma estratégia de treinamento inovadora que primeiro estabelece fluência de múltiplos patches, depois adapta para layouts mistos e, finalmente, treina um preditor separado para guiar o roteamento de tokens sem retreinar o gerador.
4. Capacidade de Integração: O framework é projetado para ser ortogonal a métodos de eficiência temporal (como cache de features), permitendo a otimização combinada.

Resultados

Experimentos foram conduzidos na geração de imagens de 1024x1024 nos benchmarks GenEval e DPG-Bench, comparando o MOSAIK contra baselines incluindo redução de patch uniforme, agendamento de patch temporal (DDiT) e cache de features (TeaCache, TaylorSeer).

• Eficiência vs. Qualidade: O MOSAIK alcança uma redução de 70% em FLOPs e uma redução de 83% no número de tokens enquanto iguala o PixelDiT de computação total no GenEval (score 0.74) e cai apenas 1.0 ponto no DPG-Bench (82.5 vs. 83.5).
• Comparação com Agendamento Temporal: Em regimes de computação agressivos (ex: redução de 60-70% de FLOPs), o MOSAIK supera consistentemente o DDiT. Enquanto o DDiT aplica um único tamanho de patch globalmente, o MOSKI REDireciona o computação dinamicamente para regiões sensíveis (ex: bordas de objetos), preservando a integridade estrutural e detalhes finos como texturas e arestas.
• Comparação com Caching: O MOSAIK supera os baselines baseados em cache (TeaCache, TaylorSeer) em orçamentos equivalentes, particularmente em regimes altamente restritos.
• Complementaridade: Quando combinado com métodos de cache, o MOSAIK sustenta o desempenho em níveis ainda mais altos de compressão (até 80% de redução de FLOPs), onde o cache sozinho causa colapso significativo de qualidade.
• Ablação: Substituir o preditor guiado por dano por um alocador aleatório ou cego ao conteúdo resulta em degradação significativa de performance, validando que o ranking de dano previsto carrega informação genuína sobre a sensibilidade espacial.

Significância e Alegações

O artigo afirma que o MOSAIK demonstra que a heterogeneidade espacial na informação visual é um fator crítico para a geração eficiente que foi negligenciado por métodos focados apenas em adaptação temporal ou de profundidade.

• Alocação Ótima de Recursos: Os autores argumentam que a criação de patches finos uniformes desperdiça computação em regiões simples, enquanto a criação de patches grosseiros uniformes destrói detalhes em regiões complexas. O MOSAIK preenche essa lacuna ao alocar recursos onde eles são mais necessários.
• Vantagem do Espaço de Pixels: O framework aproveita a redundância espacial inerente disponível em modelos de espaço de pixels (ao contrário de modelos latentes, onde a compressão do VAE já consumiu a redundância) para utilizar dinamicamente o coarse-graining sem a necessidade de otimização conjunta de ponta a ponta entre um roteador e o gerador.
• Escalabilidade: A abordagem permite um desempenho estável através de uma ampla gama de orçamentos computacionais, desde reduções moderadas até restrições agressivas, sem exigir retreinamento para diferentes níveis de orçamento.

O artigo conclui que o MOSAIK estabelece um novo padrão para eficiência em difusão de espaço de pixels, oferecendo uma alternativa altamente competitiva ao agendamento temporal e ao cache, e sugere que trabalhos futuros podem estender esses princípios para a difusão de vídeo, explorando a redundância tanto no espaço quanto no tempo.