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

The gist The BudgetPix framework proposes an adaptive tokenization method that dynamically allocates computational budget based on visual complexity and spatial layout, enabling flexible inference costs for pixel-space image diffusion models.

How it works

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. This framework seamlessly integrates with existing pixel-space diffusion architectures, enabling a single checkpoint to be operated at a wide range of compute budgets

The adaptive encoder constructs the non-uniform layout based on local intensity entropy, where cells are retained as single tokens if their Shannon entropy H(R) is below a scale-specific threshold τs + δ, and recursively split otherwise This process ensures that high-frequency details concentrate on specific regions while homogeneous areas contain minimal information The encoder utilizes a dual-path architecture where the fine path processes base patches (n = 1) directly via the pretrained patch embedder E, and the coarse path processes larger patches (n > 1) which are bilinearly downsampled and then embedded via E

The decoder reconstructs fixed-resolution images from these variable-length token sequences by mirroring the encoder's dual-path architecture The fine path maps base tokens (n=1) directly into a p×p pixel patch Pˆi, while the coarse path handles larger tokens by bilinearly upsampling and subsequently processing them through a Patch Refiner (PRs) The PRs splits the patch into 4 × 4 sub-patches, processes them through an MLP and a lightweight transformer, and fuses them with the upscaled patch to restore local texture

Training and Inference Strategy

The training objective is formalized as L(θ) = Ex,y,L,t,ϵ w(t)∥x − xˆ∥22 where w(t) represents the noise-schedule weighting derived from the base model To ensure robustness across variable token counts during training, Layout Augmentation involves three distinct layout types: Uniform layout (probability a < 1), Random layout (probability b, where a + b < 1), and Content-adaptive layout (probability 1 − a − b) The content-adaptive quadtree Q(H(x), τ + δ) is applied using the actual image entropy map H(x) during training, optimally routing tokens based on true visual complexity

During inference, the layout is estimated as the image forms, starting with a warm-up on a uniform grid to set global composition After this warm-up, the layout is recomputed every m steps (m = 5), using the entropy quadtree of the current prediction Li = Q H(ˆxi′), τ + δ∗(ˆxi′, B)

Experimental Evaluation

BudgetPix was evaluated on text-to-image generation across various architectures, including JiT, MiniT2I, PixelDiT, and MeanFlow Comprehensive assessments by human and VLM judges confirm that BudgetPix establishes a significantly improved quality-efficiency tradeoff over prior budget-adaptive baselines For instance, on 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 Quantitatively, for text-to-image generation, BudgetPix matches the fidelity of full-compute baselines while utilizing only 25% of the original token budget

Key Findings

Comprehensive experiments demonstrate that BudgetPix can be integrated seamlessly with established pixel-space architectures and provides a significantly better quality-efficiency tradeoff In class-conditional generation, BudgetPix synthesizes coherent, high-quality images using as little as 10% of the original compute budget (Figure 1, top 3 rows)–a regime where competing methods fail entirely For text-to-image generation, BudgetPix matches the fidelity of full-compute baselines while utilizing only 25% of the original token budget, achieving GenEval scores of 0.874 (vs. 0.882) on MiniT2I-L at 5122, and 0.725 (vs. 0.721) on PixelDiT at 10242 Furthermore, comprehensive evaluations by both human and Vision-Language Model (VLM) judges confirm that BudgetPix achieves a better quality-efficiency tradeoff compared to existing budget-adaptive methods In a user study, the budget can be cut by 52% for ours before human raters first see a difference, against at most 38% for the baselines The VLM-as-Judge also favors ours than every baseline below 80%

Conclusion

BudgetPix shows that a pretrained pixel-space diffusion transformer’s token count does not need to remain fixed to its pretraining configuration after one fine-tuning, serving budgets spanning from a dense grid to a small fraction of it and degrading gracefully where tokenmerging baselines collapse Its limitations are that the layout is derived from the model’s intermediate predictions, so it can miss details they do not yet show, and that full-resolution stages cap the speed-up Learning where to place tokens, not deriving it from entropy, is a natural next step

Ablation and Performance

The ablation study shows that the aggregator does most of the work in carrying content for coarse tokens, and the refiner takes roughly a further third off The performance comparison across budgets shows that BudgetPix consistently outperforms the base model across different compute budgets For class-conditional generation, BudgetPix retains FID scores of a fully dense layout even at a 60% budget, where the baselines deteriorate substantially Faithfulness across the budget indicates that ours stays closest to its dense image on nearly every metric and budget, while the base model and both baselines significantly diverge from the original image at 70% budget The wall-clock speed-up shows that BudgetPix maintains its scores across all five metrics and on the whole speed-up range, while the base model starts at the same value but collapses fast

Speed and Fidelity

In terms of wall-clock speed-up, Figure 34 plots every text-to-image metric against the measured wall-clock speed-up over the released model (DPG: DPG-Bench; IR: ImageReward) Table 6 shows that BudgetPix maintains a high speed across all budgets compared to other methods For instance, on MiniT2I-L/16 at 5122, BudgetPix achieves a wall-clock speed-up of 3.59× at the 25% budget The quality metrics, such as GenEval and DPG-Bench, show that BudgetPix achieves similar, stable scores across compute budgets while the base model degrades significantly

User Preference

The user study indicates that participants chose A in 48 to 54% of the answers that picked a side when comparing BudgetPix against RTI at all four budgets considered The VLM-as-Judge also favors ours than every baseline below 80% This confirms that BudgetPix achieves a better quality-efficiency tradeoff compared to existing budget-adaptive methods

One-Step Generation

For one-step generation with pixel MeanFlow, BudgetPix reduces the cost per step, making it complementary to inference-efficiency methods such as distillation Figure 10 shows that BudgetPix maintains FID below 4 down at the 50% budget This demonstrates that BudgetPix is still applicable to one-step generation methods by using the random layout (described in sec 3.

1. Bold token-adaptive inference capability: BudgetPix enables flexible computational budgeting at inference time by dynamically allocating compute based on visual complexity and spatial layout, allowing models to operate across variable token counts.

2. Improved resource efficiency in generation: The system achieves a significant quality-efficiency tradeoff, matching fidelity of full-compute baselines while utilizing substantially less compute, such as matching MiniT2I-L at 5122 using just 25% of the original compute budget.

3. Robustness to low-compute regimes: BudgetPix degrades gracefully under strict low-compute budgets, synthesizing coherent, high-quality images using as little as 10% of the original compute budget, a regime where competing methods fail entirely.

4. Enhanced perceptual fidelity retention: The framework demonstrates superior faithfulness across budget reductions; for instance, in JiT-B/32 at 5122, BudgetPix maintains high PSNR and SSIM even at a 70% token budget where baselines deteriorate substantially.

5. Adaptive sampling strategy: During inference, the layout is estimated as the image forms; specifically, the layout is estimated as the image forms, ensuring that the model’s own released weights run at the same budget through our sampler, without training to adapt to evolving visual data.

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 squared and PixelDiT at 1024 squared 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