💻 computer science

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

MOSAIK is a damage-guided framework for pixel-space diffusion models that dynamically allocates heterogeneous patch sizes across image regions based on estimated fidelity loss, achieving significant computational efficiency (70% FLOPs reduction) while maintaining competitive generation quality compared to full-compute baselines.

Original authors: Mohammadreza Hami, Mohammadreza Samadi, Chao Gao, Negar Hassanpour

Published 2026-08-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Mohammadreza Hami, Mohammadreza Samadi, Chao Gao, Negar Hassanpour

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to paint a masterpiece on a giant canvas, but you only have a limited supply of paint and a very strict time limit. In the world of artificial intelligence, creating images is a bit like this. For a long time, computers used a "compressed" method to draw pictures, kind of like sketching a scene in a tiny notebook before trying to fill in the details. This was fast, but it had a ceiling: once you sketched it small, you couldn't magically recover the tiny, sharp details like the texture of fur or the individual petals of a flower.

To get around this, a newer generation of AI started painting directly on the giant canvas, pixel by pixel. This is like switching from a tiny notebook to a massive mural. The result? Stunningly realistic images. But there's a catch: painting every single pixel requires the computer to process a massive amount of information, which is slow and expensive. To speed things up, scientists have tried using "big brushes" to cover large areas of the canvas with a single stroke. However, using a big brush everywhere is a bit like trying to paint a portrait of a tiger with a house-painting roller; you might get the shape right, but you'll lose the stripes and the whiskers. The big question has been: How do we use big brushes for the boring background (like the sky) but tiny, precise brushes for the important parts (like the tiger's face) without the computer getting confused?

This is where a new method called MOSAIK comes in. Think of MOSAIK as a super-smart art director for the AI painter. Instead of forcing the computer to use the same brush size for the entire image, MOSAIK looks at the picture as it's being created and decides, "Hey, this part is just a blurry sky, let's use a giant brush here to save time. But this part is the tiger's eye, let's switch to a tiny, super-precise brush right now."

The paper introduces MOSAIK (which stands for Multi-Patch Content-Aware Spatial Allocation of Image Tokens) as a way to make these high-quality, pixel-by-pixel image generators much faster without ruining the picture. The researchers found that by dynamically changing the size of the "brush strokes" (or patches) across different parts of the image, they could cut the computer's workload by a huge amount. Specifically, they managed to reduce the number of calculations (called FLOPs) by 70% and the number of data chunks (tokens) by 83%.

Here is the magic trick: MOSAIK doesn't just guess where to use big brushes. It uses a "damage predictor." Imagine the AI has a special sensor that asks, "If I use a big brush on this specific spot, how much detail will I lose?" If the answer is "a lot" (like on a tiger's fur), it keeps the brush small. If the answer is "nothing" (like on a smooth blue sky), it switches to a big brush. The paper shows that this smart, uneven approach works much better than older methods that tried to speed things up by just using big brushes for the whole image at certain times or by skipping steps.

In their tests, MOSAIK was able to generate images that looked almost exactly the same as the slow, full-power version, even though it was doing 70% less work. When they compared it to other speed-up tricks, MOSAIK kept the details sharp while the others started to look blurry or lost important features. The researchers suggest that this "damage-guided" approach is a powerful way to make high-quality AI image generation faster and more efficient, proving that you don't have to choose between speed and quality if you know where to be efficient and where to be precise.

Technical Summary: MOSAIK

Problem Statement

Pixel-space diffusion models (PDMs) offer a significant advantage over latent diffusion models (LDMs) by avoiding the reconstruction ceiling imposed by Variational Autoencoders (VAEs), thereby preserving fine-grained, high-frequency spatial details. However, generating directly in pixel space results in a substantially higher token count. Since self-attention mechanisms in transformers exhibit quadratic complexity with respect to token count, PDMs incur prohibitive computational costs.

Existing efficiency methods attempt to mitigate this by varying patch sizes (the spatial area represented by a single token) across denoising timesteps or network depth. While these approaches reduce the total token count, they apply a spatially uniform patch layout across the entire image at any given step. This uniformity creates a fundamental inefficiency: it either overspends computation on simple, low-information regions or, conversely, constrains representational capacity in sensitive regions where fine details are critical. The paper argues that visual information density is inherently heterogeneous, necessitating a dynamic, per-region allocation strategy that current methods lack.

Methodology: MOSAIK

MOSAIK (Multi-patch cOntent-aware Spatial Allocation of Image toKens) is a framework designed to vary patch sizes across different image regions and denoising steps based on the predicted "damage" (fidelity loss) caused by coarsening. The approach is built upon the PixelDiT backbone and utilizes a three-stage training pipeline:

1. Stage 1: Multi-Patch Fluency

The goal is to equip the model to operate at multiple uniform patch sizes ( $P \in \{16, 32, 64\}$ ) without catastrophic forgetting of the original generative prior.

Architecture: The backbone weights are frozen. Low-Rank Adaptation (LoRA) modules are injected into the DiT blocks. Separate tokenizers are introduced for $p32$ and $p64$ , initialized by composing the pretrained $p16$ projection with bilinear downsampling.
Training: The model undergoes self-distillation, optimizing to match the velocity predictions of the frozen $p16$ teacher model across all uniform patch sizes.

2. Stage 2: Mixed-Patch Adaptation

This stage enables the model to process arbitrary, heterogeneous patch layouts (where different regions use different patch sizes) within a single forward pass.

Positional Embeddings: Standard Rotary Positional Embeddings (RoPE) assume a uniform grid. MOSAIK addresses the spatial mismatch of mixed patches by anchoring each token to the geometric center of its patch within a normalized $p16$ coordinate frame.
Training: The generator is trained on randomly sampled heterogeneous layouts, deliberately decoupling the spatial arrangement of patch sizes from the image content. This forces the model to learn to execute any layout regardless of its alignment with visual features.

3. Stage 3: Damage-Guided Layout Predictor

This stage introduces a lightweight predictor that dynamically allocates tokens based on the estimated degradation caused by coarsening specific regions.

Damage Definition: "Damage" is defined as the squared difference between the prediction of a uniform $p64$ layout and the high-fidelity $p16$ teacher, averaged over non-overlapping $64 \times 64$ regions.
Predictor Training: A lightweight convolutional network with FiLM conditioning is trained to regress this damage map directly from intermediate backbone features (specifically DiT block-6 activations) and text embeddings. This avoids the need for a secondary forward pass during inference.
Inference Strategy: At each denoising step, the predictor estimates the damage map. Regions with high predicted damage are assigned finer patches ( $p16$ ), while low-damage regions are assigned coarser patches ( $p32$ or $p64$ ). The allocation is constrained by a user-specified token budget, prioritizing refinements that offer the highest damage reduction per added token.

Key Contributions

Spatially Adaptive Patchification: MOSAIK is the first framework to apply heterogeneous patch layouts within a pixel-space diffusion model, moving beyond the spatial uniformity constraints of prior efficiency methods.
Damage-Guided Allocation: The authors formulate patch allocation as a damage-prediction problem, using a lightweight predictor to identify regions where coarsening would cause the most significant fidelity loss.
Three-Stage Training Pipeline: A novel training strategy that first establishes multi-patch fluency, then adapts to mixed layouts, and finally trains a separate predictor to guide token routing without retraining the generator.
Integration Capability: The framework is designed to be orthogonal to temporal efficiency methods (like feature caching), allowing for combined optimization.

Results

Experiments were conducted on 1024x1024 image generation using the GenEval and DPG-Bench benchmarks, comparing MOSAIK against baselines including uniform patch reduction, temporal patch scheduling (DDiT), and feature caching (TeaCache, TaylorSeer).

Efficiency vs. Quality: MOSAIK achieves a 70% reduction in FLOPs and an 83% reduction in token count while matching the full-compute PixelDiT on GenEval (0.74 score) and dropping only 1.0 point on DPG-Bench (82.5 vs. 83.5).
Comparison with Temporal Scheduling: In aggressive compute regimes (e.g., 60-70% FLOPs reduction), MOSAIK consistently outperforms DDiT. While DDiT applies a single patch size globally, MOSAIK dynamically redirects compute to sensitive regions (e.g., subject boundaries), preserving structural integrity and fine details like textures and edges.
Comparison with Caching: MOSAIK outperforms caching-based baselines (TeaCache, TaylorSeer) at matched budgets, particularly in highly constrained regimes.
Complementarity: When combined with caching methods, MOSAIK sustains performance at even higher compression levels (up to 80% FLOPs reduction), where caching alone causes significant quality collapse.
Ablation: Replacing the damage-guided predictor with a random or content-blind allocator results in significant performance degradation, validating that the predicted damage ranking carries genuine information about spatial sensitivity.

Significance and Claims

The paper claims that MOSAIK demonstrates that spatial heterogeneity in visual information is a critical factor for efficient generation that has been overlooked by methods focusing solely on temporal or depth-wise adaptation.

Optimal Resource Allocation: The authors argue that uniform fine patchification wastes computation on simple regions, while uniform coarse patchification destroys details in complex regions. MOSAIK bridges this gap by allocating resources where they are most needed.
Pixel-Space Advantage: The framework leverages the inherent spatial redundancy available in pixel-space models (unlike latent models where VAE compression has already consumed redundancy) to dynamically utilize coarsening without the need for end-to-end joint optimization of a router and generator.
Scalability: The approach allows for stable performance across a wide range of compute budgets, from moderate reductions to aggressive constraints, without requiring retraining for different budget levels.

The paper concludes that MOSAIK sets a new standard for efficiency in pixel-space diffusion, offering a highly competitive alternative to temporal scheduling and caching, and suggests that future work could extend these principles to video diffusion, exploiting redundancy across both space and time.