BudgetPix: Compute-Adaptive Tokenization for Pixel-Space Image Diffusion

AI summary

Overview

Research area: Computer vision and generative modeling — specifically efficient, pixel-space diffusion transformers for image generation.

Technical level: Intermediate. The paper assumes familiarity with diffusion/flow-matching models, transformers, and token-based image representations, but its central idea (spending more compute on detailed regions, less on plain ones) is intuitive without that background.

Scope: BudgetPix is a tokenization framework that lets a single pretrained pixel-space diffusion transformer run at variable compute budgets at inference time by encoding images into variable-length, multi-scale token sequences driven by local visual entropy.

What This Paper Is About

Pixel-space diffusion transformers divide an image into a fixed grid of equally sized patches, giving every patch the same computational budget — a plain sky costs as much as an intricate texture, and the total cost is locked in at training time. Because self-attention scales quadratically with token count, this uniform scheme is expensive and cannot be adjusted for slower or faster hardware at deployment. The paper's goal is to make tokenization content-adaptive, so a single fine-tuned checkpoint can spend fewer tokens on simple regions and more on complex ones, and can be dialed up or down to any budget without retraining.

Key Contributions

1. A compute-adaptive tokenization framework (BudgetPix) that plugs into existing pixel-space diffusion architectures — the paper demonstrates it on JiT, MiniT2I, PixelDiT, and MeanFlow — and converts a fixed-budget checkpoint into one that operates across a range of budgets.
2. An entropy-guided encoder that builds a non-uniform quadtree layout from local intensity histograms and encodes the resulting variable-sized patches into same-dimension tokens using a multi-scale patch embedder with a dedicated coarse path (the Patch Aggregator).
3. A scale-aware decoder that reconstructs fixed-resolution images from multi-scale tokens, using a lightweight Patch Refiner (two self-attention layers with AdaLN conditioning) to re-inject high-frequency detail lost when coarse patches are upsampled; it adds roughly 1% to the parameter count.
4. A training and sampling schedule for variable token counts, including layout augmentation (uniform, random, and content-adaptive layouts mixed during training) and an inference procedure that runs warm-up steps on the uniform grid, then estimates the layout from the model's own current prediction and rebuilds it every m steps.

Main Findings

• Text-to-image at a quarter of the budget: BudgetPix matches the fidelity of MiniT2I-L at 512² and PixelDiT at 1024² using just 25% of the original compute budget, with GenEval scores of 0.874 (vs. 0.882) on MiniT2I-L and 0.725 (vs. 0.721) on PixelDiT.
• Class-conditional with near-zero degradation: On a MeanFlow backbone, BudgetPix requires merely 60% of the full compute budget for near-zero quality loss, with a marginal FID increase of 0.81 points (from 2.65 to 3.46).
• Graceful degradation where baselines collapse: On JiT-L/32 at 512², BudgetPix moves from FID 2.71 at 100% budget to 4.41 at 60%, 5.57 at 50%, and 11.54 at 25%, while the base model goes from 2.66 to 107.51 and ToMe/FeatSim fare worse still. On class-conditional generation, the method produces coherent images at as little as 10% of the original budget.
• PixelDiT stability across budgets: GenEval for the base model falls from 0.721 at 100% to 0.323 at 25%, while BudgetPix holds 0.741 at 100% and 0.725 at 25%; DPG-Bench falls from 84.8 to 55.1 for the base model versus 85.2 to 83.7 for BudgetPix.
• Better faithfulness to the dense image: Measured against each method's own full-budget output (paired PSNR, SSIM, LPIPS, CLIP and DINOv2 similarity), BudgetPix stays closest on nearly every metric and budget, while the base model and merging baselines diverge noticeably by 70% budget.
• Human and VLM judges agree: In a user study with 50 people per form over 20 GenEval prompts, and a VLM-as-Judge setup over 100 prompts, the budget can be cut by 52% for BudgetPix before human raters first see a difference, against at most 38% for baselines; the VLM judge tolerates 36% versus at most 21%. Human raters find BudgetPix more faithful than RTI at all four budgets considered, and the VLM favors it over every baseline below 80%.
• Works with one-step generation: On one-step (1-NFE) pixel MeanFlow pMF-L/16 at 256², the released weights collapse at 80% budget, while BudgetPix keeps FID below 4 down to a 50% budget.
• Ablation: On JiT-B at 512², adding the Patch Aggregator drops the 25%-budget FID from 190.71 (base) to 25.07, and adding the Patch Refiner further to 16.71. The aggregator does most of the work; the refiner contributes roughly another third and matters more as the budget falls.

Methodology in Plain English

The starting observation is that images are not uniformly informative: a clear sky carries little information, while a textured object carries a lot. BudgetPix exploits this by deciding, per image, which regions get a small patch (one token) and which get merged into a bigger patch (also one token, but covering more pixels).

To decide, the method divides the image into a coarse grid of cells, builds a histogram of pixel intensities inside each cell, and measures its Shannon entropy. Low-entropy cells stay as a single token. High-entropy cells are split into four quadrants and re-evaluated, recursively, until the finest patch size is reached. The result is a tiling of the image with a mix of patch sizes and a token count that depends on content.

Because the tokenizer must handle patches of different sizes, the encoder has two paths. Small patches go through the base model's existing patch embedder. Large patches are downsampled to the small size and embedded the same way, but a side branch (the Patch Aggregator) breaks the large patch into sub-patches, embeds each, and mixes them into a correction that restores the lost fine detail. A learned scale embedding tells the model how big each token's patch was.

The decoder mirrors this. Small tokens go through the original decoding head. Large tokens are decoded to a small patch, upsampled to their true size, and refined by a small transformer (the Patch Refiner) that adds back the missing high frequencies as a residual.

The tricky part is generation: at inference there is no finished image to compute entropy from. The model first runs a few warm-up denoising steps on the ordinary uniform grid to get a rough predicted image, computes the layout from that prediction, and refreshes the layout every m steps. During fine-tuning, the layout is sometimes the true content-adaptive one, sometimes the standard uniform grid (to stay anchored to the pretrained weights), and sometimes a randomized layout, so the model learns to handle coarse tokens even when they happen to cover detailed regions. The global threshold offset is randomized per image, which teaches the model to work anywhere from roughly 25% to 100% of the dense grid. The whole pipeline is fine-tuned with the base model's own flow-matching objective, and the new modules are zero-initialized so training begins exactly at the released checkpoint.

Why This Matters

The work reframes a hardcoded architectural choice — the token grid — as a runtime control knob. For pixel-space diffusion, where attention cost grows quadratically with tokens, that means one trained model can serve many latency and hardware targets instead of one model per budget. It also offers an alternative to token merging for generative models, which the paper shows degrades sharply at low budgets, and it is complementary to step-reduction techniques like distillation, since BudgetPix cuts cost per step rather than the number of steps.

Real-world applications implied by the framing:

• On-device and edge image generation, where memory and compute are limited and the paper explicitly motivates deployment on resource-constrained devices.
• Latency-tiered creative tools, where an interactive interface can trade quality for responsiveness depending on user demand.
• Batch or high-volume generation pipelines, where a fixed serving budget can be met without retraining a separate model.
• Bandwidth- or cost-constrained API services, where a single checkpoint covers multiple pricing or quality tiers.

Industry relevance: the method is designed as a drop-in wrapper around released checkpoints rather than a new architecture, and it is demonstrated on several public model families plus a one-step MeanFlow model — a practical profile for teams that already have a trained pixel-space model and want variable-cost inference without starting over.

Future Directions

• Learning the layout instead of deriving it from entropy. The authors name this explicitly as a natural next step, replacing the hand-designed quadtree rule with a learned placement policy.
• Handling details the model has not yet imagined. The paper's stated limitation is that the layout comes from the model's intermediate predictions, so it can miss details those predictions do not yet show — an open question for how to anticipate complexity earlier in sampling.
• Removing the ceiling imposed by full-resolution stages. The authors note that stages such as PixelDiT's pixel-level second stage cap the achievable speed-up, implying further work on making those stages budget-aware.
• Extending the idea beyond the token grid. Since BudgetPix reduces per-step cost, combining it with step-reduction or distillation methods — including the one-step setting the paper touches on — is a natural direction for compounding efficiency gains.

Target Audience

Researchers and engineers working on efficient generative models, particularly diffusion transformers and pixel-space architectures; practitioners who want variable-cost inference from an existing checkpoint; and readers interested in content-adaptive tokenization, since the paper bridges that literature from discriminative vision into generative denoising. It is also useful for those evaluating budget-adaptive baselines, given its paired faithfulness metrics and human/VLM judging protocols.

Authors’ abstract

Most image generation models rely on uniform tokenization, allocating the exact same computational budget to equally-sized image patches. This static paradigm cannot adapt to different resource constraints at inference time, and yields suboptimal quality-cost tradeoff by devoting the same effort to both plain backgrounds and intricate details. We propose BudgetPix, an adaptive tokenization framework that dynamically allocates compute based on visual complexity and spatial layout, enabling flexible computational budgeting at inference time. BudgetPix comprises three key components: (1) an adaptive encoder that maps a fixed-size image to a variable-length token sequence using an entropy-guided quadtree alongside a multi-scale patch embedder; (2) a scale-aware decoder reconstructs fixed-resolution images from multi-scale token sets; and (3) a flexible training and sampling schedule that enables pixel-space denoisers to operate across variable token counts. BudgetPix seamlessly integrates with existing pixel-space diffusion architectures, enabling a single checkpoint to be operated at a wide range of compute budgets. Evaluated on text-to-image generation, BudgetPix matches the fidelity of MiniT2I-L at $512^2$ and PixelDiT at $1024^2$ using just 25% of the original compute budget. In class-conditional generation using a MeanFlow backbone, BudgetPix requires merely 60% of the full compute budget to produce images with near-zero quality degradation, observing a marginal 0.8-point increase in FID. Comprehensive assessments by human and VLM judges confirm that BudgetPix establishes a significantly improved quality-efficiency tradeoff over prior budget-adaptive baselines. More details are available at our project page: https://karaozgur.com/BudgetPix