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

2026-08-05

AI Summary

This work addresses the high computational cost of pixel-level diffusion models, which stems from the large number of image tokens required for self-attention. While existing approaches reduce token count using uniform patch sizes, they overlook regional variations in sensitivity to spatial coarsening. To overcome this limitation, the authors propose MOSAIK, a framework that enables heterogeneous patch layouts both across image regions and denoising timesteps. A lightweight predictor estimates fidelity loss induced by coarsening in each region, allowing dynamic allocation of fine-grained representations under a fixed token budget. Built upon a modified PixelDiT architecture that supports arbitrary heterogeneous patches, MOSAIK trains its damage predictor using intermediate denoising features to achieve content-aware token allocation. With 70% fewer FLOPs and 83% fewer tokens, MOSAIK matches the GenEval performance of full-resolution PixelDiT and incurs only a 1.0-point drop on DPG-Bench, significantly outperforming current efficient generative methods.

Abstract

Pixel-space diffusion models avoid the reconstruction ceiling of latent diffusion models by generating directly in image space. However, their substantially higher token count makes generation expensive due to the quadratic complexity of self-attention. Several existing efficiency methods reduce this cost by using larger patches at selected denoising steps, thereby representing the image with fewer tokens. Yet, each step still uses a single patch size uniformly across the entire image, overlooking that different regions suffer different fidelity losses when coarsened. We introduce MOSAIK, a damage-guided framework that varies patch size across regions and denoising steps. MOSAIK adapts the PixelDiT backbone to generate arbitrary heterogeneous patch layouts, and a lightweight predictor uses intermediate denoising features to estimate the fidelity loss caused by coarsening each region. Given a token budget, our damage-guided layout predictor assigns fine patches to sensitive regions and coarse patches elsewhere. Remarkably, while reducing FLOPs by 70% and token count by 83%, MOSAIK matches the full-compute PixelDiT on GenEval and its DPG-Bench score drops by only 1.0 point. Compared to diverse efficiency paradigms, including temporal patch scheduling and feature caching, our approach delivers highly competitive performance at moderate budgets and consistently outperforms these baselines in highly constrained compute regimes.

Problem

• pixel-space diffusion
• token efficiency
• patch size allocation
• fidelity loss
• self-attention complexity

Innovation

• content-aware patch allocation
• damage-guided token layout
• pixel-space diffusion
• heterogeneous patch scheduling
• efficient image generation