PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ozgur Kara, Yujia Chen, Daniel Watson, David Forsyth, James Matthew Rehg, Wen-Sheng Chu, Du Tran

arXiv:2610.12307v1cs.CV

TL;DR

Pixel-space diffusion models face high costs from uniform tokenization and fixed inference budgets. BudgetPix adapts tokenization and training to support variable token counts; it reports strong quality-efficiency results across text-to-image and class-conditional generation.

  • Problem

    Uniform tokenization makes pixel-space diffusion costly and fixes inference compute, limiting adjustment to resource constraints and latency needs.

    Source: §1 · p. 2

  • Method

    BudgetPix adapts pretrained pixel-space diffusion transformers with an entropy-driven variable-length token encoder, a fixed-resolution decoder, and training and sampling across variable token counts.

    Source: §1 · p. 3

  • Results

    BudgetPix improves quality-efficiency across tasks: text-to-image matches full-compute fidelity at 25% of the original token budget, while MeanFlow class-conditional generation at 60% sees a 0.81-point FID increase, from 2.65 to 3.46.

    Source: §1 · p. 3

  • Takeaways & Limitations

    After one fine-tuning, a single checkpoint serves budgets from a dense grid to a small fraction of it, degrading gracefully where token-merging baselines collapse.

    Source: §5 · p. 10

    The layout comes from intermediate predictions and can miss details they do not yet show; full-resolution stages such as PixelDiT’s pixel stage cap speed-up.

    Source: §5 · p. 10

Abstract

from arXiv · show

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

1 INTRODUCTION

Uniform tokenization makes pixel-space diffusion expensive and fixes inference cost, despite uneven visual complexity. BudgetPix adapts token allocation and reports strong quality-efficiency results across generation settings.

1 INTRODUCTION

  • Uniform tokenization spends equal compute on low-entropy backgrounds and detailed regions, while pixel-space models face quadratic attention costs and rigid inference budgets.For a 512 × 512 flag scene, the figure contrasts 1,024 uniform tokens with 307 BudgetPix tokens; for a textured image, it contrasts 256 with 127.Source: §1 · pp. 2–3, Figure 2 · p. 2
  • Generative token allocation is difficult because the denoiser must estimate spatial complexity before the image content has fully materialized.This differs from discriminative models, which allocate tokens over an already-provided image.Source: §1 · p. 3
  • BudgetPix adapts pretrained pixel-space diffusion transformers through an entropy-driven variable-length encoder, a fixed-resolution decoder, and training and inference strategies for variable token counts.The framework is reported to integrate with JiT, MiniT2I, PixelDiT, and MeanFlow.Source: §1 · p. 3
  • At just 10% of the original compute budget, BudgetPix produces coherent, high-quality class-conditional images in a regime where competing methods fail.The authors report graceful degradation under strict low-compute budgets across class-conditional and text-to-image models from 256^2 to 1024^2.Source: §1 · p. 3
  • At 25% of the original token budget, BudgetPix matches full-compute fidelity, with GenEval 0.874 vs. 0.882 on MiniT2I-L at 512^2 and 0.725 vs. 0.721 on PixelDiT at 1024^2.Human and VLM evaluations also find a better quality-efficiency tradeoff than existing budget-adaptive methods.Source: §1 · p. 3

2 RELATED WORK

Pixel-space generation evolved toward plain-transformer models, but their uniform patch grids keep compute tied to patch size. Prior adaptive tokenization methods largely address discriminative vision, while generation must choose a layout before the image exists.

2 RELATED WORK

  • Pixel-space models developed from U-Nets and cascaded designs through resolution-shifting schedules and hierarchical transformers, then JiT and PixelDiT used plain transformers to match latent-model quality at equal sequence length.Latent diffusion uses a pretrained autoencoder, with a reconstruction ceiling and a second training stage; pixel-space approaches bypass the autoencoder.Source: §2 · pp. 3–4
  • JiT and PixelDiT use a single uniform patch grid, so patch size fixes computational cost; PixelDiT adds a pixel-level stage for large-patch decoding, while PixNerd uses a per-patch neural field.The cited work describes these as responses to the weakness of decoding large patches.Source: §2 · p. 4
  • Prior content-adaptive tokenization focused on discriminative vision, while generative methods vary updates or caching, vary patch size uniformly, retrain latent tokenizers, or adapt region tokens with LoRA.BudgetPix instead uses its adaptive encoder and decoder, fully fine-tunes the checkpoint, and derives layout from the model’s clean-image estimate.Source: §2 · p. 4

3 BUDGETPIX

BudgetPix replaces uniform tokenization with a content-adaptive pipeline that lets pixel-space diffusion backbones process variable token counts. Its encoder, decoder, and training and inference strategy support multi-scale patches and flexible budgets.

3.1 OVERVIEW

  • BudgetPix combines an encoder, decoder, and unified training and inference strategy to let existing diffusion backbones adjust token counts to visual complexity.The overview describes the encoder and decoder as trainable modules around existing pixel-space diffusion transformers.Source: §3.1 · p. 4, Figure 3 · p. 4

3.2 BUDGETPIX ENCODER: EMBED NON-UNIFORM PATCHES

  • Entropy-guided quadtree splitting retains a cell as one token below a scale-specific entropy threshold and otherwise recursively divides it into four finer cells.Splitting continues to the finest p × p scale, so the resulting multi-scale tokens tile the image.Source: §3.2 · p. 5, Figure 4 · p. 5
  • The encoder embeds fine patches directly and processes coarse patches through downsampling and localized sub-patch aggregation before fusing them into tokens.The coarse path adds scale-specific information to the base embedding to recover details lost by downsampling.Source: §3.2 · p. 5

3.3 BUDGETPIX DECODER: DECODING MULTI-SCALE IMAGE PATCHES

  • The decoder upsamples coarse outputs and uses a local refiner to add high-frequency detail, while fine tokens use the pretrained decoding head directly.The refiner operates within each coarse token using short local sequences and adds about 1% to the parameter count.Source: §3.3 · pp. 5–6

3.4 TRAINING AND INFERENCE

  • Training mixes uniform, random, and content-adaptive layouts, while inference warms up on a uniform grid before rebuilding layouts from predicted-image entropy.Random layouts expose coarse tokens to detailed regions; randomized thresholds train across roughly 25%–100% of the dense-grid token count, and inference rebuilds layouts every m steps.Source: §3.4 · pp. 6–7

4 EXPERIMENTS

Experiments fine-tune released pixel-space diffusion models for variable token layouts and compare them with budget-matched baselines. Across class-conditional and text-to-image settings, BudgetPix preserves quality more effectively as compute falls, with additional tests of faithfulness, human preference, speed, and component contributions.

4.1 IMPLEMENTATION DETAILS / 4.2 EVALUATION SETTINGS

  • Seven released pixel-space transformers, including one-step pMF-L/16, are fine-tuned once with layout augmentation; comparisons hold prompts or labels, noise, solver, and step count fixed while varying budget and reduction rule.Evaluation uses FID-50k and Inception Score for class-conditional quality, GenEval and DPG-Bench for prompt adherence, preference metrics, faithfulness measures, and wall-clock cost.Source: §4.1 · p. 7, §4.2 · pp. 7–8

4.3 ANALYSIS

  • At 50–60% budgets, JiT-L/32 merging baselines have FID several times higher than BudgetPix, which retains fully dense-layout FID at 60%.The base model and merging baselines' FID rise sharply as the budget decreases; Inception Score follows the same trend.Source: §4.3 · p. 8, Figure 8 · p. 8
  • At 25% of the original token budget, BudgetPix maintains stable GenEval and DPG-Bench scores on MiniT2I-L/16 and similar performance trends on PixelDiT-T2I at 1024².The supplied PixelDiT table caption reports that BudgetPix consistently outperforms the base model across compute budgets.Source: Table 1 · p. 8, Figure 7 · p. 8, §4.3 · p. 9
  • Human raters first notice a difference after a 52% compute cut for BudgetPix versus at most 38% for baselines; VLM judges' thresholds are 36% versus at most 21%.BudgetPix is also preferred to RTI by human raters at all four tested budgets and favored over every baseline by the VLM judge below 80%.Source: §4.3 · p. 9, Figure 9 · p. 9
  • BudgetPix maintains text-to-image scores across the measured speed-up range, while merging baselines gain less speed because merging itself takes time.On one-step pMF-L/16, Figure 10 reports a better quality-efficiency trade-off than the base model and other token-merging baselines.Source: Figure 7 · p. 8, §4.3 · p. 9, Figure 10 · p. 9
  • The Patch Aggregator contributes most in the ablation, while the Refiner adds roughly a further third and matters more at lower budgets.The supplied ablation caption says the combination of Patch Aggregator and Refiner works best across compute budgets.Source: Table 2 · p. 9, §4.3 · p. 10

5 CONCLUSION

BudgetPix lets one fine-tuned pixel-space diffusion checkpoint serve token budgets from a dense grid to a small fraction of it. The conclusion also identifies limits in its prediction-derived layouts and in speed-ups capped by full-resolution stages.

5 CONCLUSION

  • A single fine-tuned checkpoint can serve budgets ranging from a dense grid to a small fraction of it, degrading gracefully where token-merging baselines collapse, even at 1 NFE.The conclusion frames this as evidence that pretrained pixel-space diffusion transformers need not retain their original fixed token count.Source: §5 · p. 10
  • Layouts derived from intermediate predictions can miss details not yet visible, while full-resolution stages such as PixelDiT’s pixel stage cap speed-up.The conclusion names learning token placement rather than deriving it from entropy as a natural next step.Source: §5 · p. 10

A ADDITIONAL QUALITATIVE RESULTS ACROSS THE BUDGET

The appendix supplies qualitative comparisons and samples across models, prompts or classes, token budgets, and—where indicated—visualized token layouts.

A.1 BASELINE COMPARISON (MINIT2I-L/16, 5122)

  • Baseline comparisons on MiniT2I-L/16 at 512² arrange methods by row and budgets from 100% to 10% by column across prompts.Figures 11–16 compare the base model, ToMe, FeatSim, RTI, and BudgetPix.Source: §A.1 · p. 16, Figure 11 · p. 16

A.2 BUDGETPIX (JIT-L/32, 5122)

  • JiT-L/32 at 512² is shown with additional samples organized by ImageNet class and token budget from 100% to 10%.Figures 17 and 18 provide the qualitative results.Source: §A.2 · p. 19, Figure 17 · p. 19

A.3 BUDGETPIX (MINIT2I-L/16, 5122)

  • MiniT2I-L/16 at 512² is shown with additional samples organized by prompt and token budget from 100% to 10%.Figures 19 and 20 provide the qualitative results.Source: §A.3 · p. 21, Figure 19 · p. 21

A.4 BUDGETPIX (PIXELDIT-T2I, 10242)

  • PixelDiT-T2I at 1024² is shown with additional prompt-based samples across token budgets from 100% to 10%.Figures 21 and 22 provide the qualitative results.Source: §A.4 · p. 23, Figure 21 · p. 23

A.5 BUDGETPIX (PMF-L/16, 2562, 1 NFE)

  • One-step pMF-L/16 samples at 256² cover ImageNet classes and budgets from 100% to 25%; Figure 23 also includes the released model at 100% and an untrained-layout base model at 50%.Figures 23–25 show the additional qualitative results.Source: §A.5 · p. 25, Figure 23 · p. 25

A.6 BUDGETPIX LAYOUTS (JIT-L/32, 5122)

  • JiT-L/32 layout figures pair samples at 100%, 50%, and 25% of the grid with the same images overlaid by token outlines colored by size.Rows correspond to ImageNet classes.Source: §A.6 · p. 28, Figure 26 · p. 28

A.7 BUDGETPIX LAYOUTS (MINIT2I-L/16, 5122)

  • MiniT2I-L/16 layout figures pair prompt-based samples at 100%, 50%, and 25% of the grid with images overlaid by token outlines colored by size.Rows correspond to prompts.Source: §A.7 · p. 30, Figure 28 · p. 30

A.8 BUDGETPIX LAYOUTS (PIXELDIT-T2I, 10242)

  • PixelDiT-T2I layout figures pair prompt-based samples at 100%, 50%, and 25% of the grid with images overlaid by token outlines colored by size.Rows correspond to prompts.Source: §A.8 · p. 32, Figure 30 · p. 32

B RESULTS IN FULL

This block reports generation, fidelity, speed, and budget-mode results, then documents the added modules, training objective, checks, and fine-tuning settings. The supplied passages describe the scope and protocols of several results but do not provide most table values.

B.1 CLASS-CONDITIONAL GENERATION (JIT)

B.3 TEXT-TO-IMAGE GENERATION (MINIT2I)

  • Table 5 reports text-to-image values by token budget for MiniT2I-B/16 and MiniT2I-L/16 at 512^2; PixelDiT at 1024^2 is reported in Table 1.The supplied passage identifies the tables and settings but not their cell values.Source: §B.3 · p. 36, Table 5 · p. 36

B.4 WALL-CLOCK SPEED-UP (EVERY PARENT)

  • Wall-clock speed-up is measured against each released model at its dense grid, with figures and tables covering speed-up, quality, and seconds per image across parents.PixelDiT timing covers only the base model and BudgetPix.Source: §B.4 · p. 37, Figure 34 · p. 37, Figure 35 · p. 37

B.5 BUDGET MODE AGAINST THRESHOLD MODE (JIT-L/32)

  • At the same mean token count, budget and threshold modes lie on nearly one curve, but only budget mode guarantees each image stays within its requested token count.The paper therefore reports its other results in budget mode.Source: §B.5 · p. 41, Figure 37 · p. 41, Table 14 · p. 41

B.6 ONE-STEP GENERATION (PIXEL MEANFLOW)

  • The one-step pMF-L/16 ImageNet experiment uses an entropy-quadtree layout borrowed from one of 20 same-class training images and reports FID-50k and Inception score.The passage specifies 256^2 resolution and a 256-token dense grid; guidance follows the released model at every budget.Source: §B.6 · p. 42, Table 15 · p. 42

C.1 ARCHITECTURE OF THE ADDED MODULES

  • Added coarse-token modules embed and decode patches at scales n = 2 and 4; CellRefine adjusts each cell’s conditioning of PixelDiT’s pixel stage, whose per-pixel decoding sets a cost floor.Base tokens retain the released embedding and head, and each coarse scale has its own Patch Aggregator.Source: §C.1 · p. 43, Table 16 · p. 43

C.2 TRAINING OBJECTIVE

  • The training objective preserves each parent’s pretrained weighting, while the velocity-regression description expresses the target and prediction through the clean-image estimate.PixelDiT uses the same objective with w = σ^-2 and its own logit-normal noise sampling.Source: §C.2 · p. 44

C.3 ZERO-INITIALISATION CHECK

  • The zero-initialisation check compares added-path outputs with released-model outputs on identical inputs; the shared recipe specifies entropy-quadtree layouts, training augmentation, and layout recomputation every five sampling steps.Every parent starts from its released checkpoint, and one checkpoint per parent serves every budget; the pMF parent uses its own objective and fixed-per-step layouts.Source: §C.3 · p. 45, Table 17 · p. 45, §C.4 · pp. 46–47

C.5 EVALUATION PROTOCOL

This section defines the comparison methods and evaluation measures for class-conditional and text-to-image generation, dense-image faithfulness, and human- or model-based judgments.

Methods compared

  • Comparisons use the released model at its dense grid and a base model run at the same budget without training; ToMe and FeatSim wrap released weights, while RTI is available only for MiniT2I-B and MiniT2I-L.Within each parent, methods share prompts or labels, noise, solver, steps, guidance, and attention arithmetic.Source: §C.5 · p. 48

Class-conditional

  • Class-conditional FID and Inception score use native-resolution uint8 images and ImageNet-1k training-set statistics; each FID is a single run.FID-50k uses 50 images per class, while ablations use FID-10k with 10 per class.Source: §C.5 · p. 48

Text-to-image

  • Text-to-image evaluation uses GenEval, DPG-Bench, CLIP score, PickScore, and ImageReward; dense-grid held-out FID compares generated and real CC12M or PD12M images.GenEval requires every prompt element to be rendered; the other reported measures use their specified prompt sets and scoring procedures.Source: §C.5 · p. 48

Paired metrics

  • Paired fidelity scores compare each reduced-budget image with the same method’s dense-grid image generated from the same prompt or label and noise.Measures include PSNR, SSIM, LPIPS, and cosine similarity of CLIP and DINOv2 embeddings; JiT PSNR is also split by merged versus retained pixels.Source: §C.5 · p. 48

C.6 USER STUDY AND VLM-AS-JUDGE

The section describes user and VLM judging protocols, wall-clock measurement checks, and development-model decoder and sampler ablations. Its supplied results include position-bias and participant-choice checks but omit most ablation values.

C.6 USER STUDY AND VLM-AS-JUDGE

  • The user study asks when a visible difference first appears across descending budgets and which candidate is more faithful; the VLM sees a dense-grid reference and reduced-budget candidates.The studies use shared prompts and noise across compared methods and budgets, with left-right order counterbalanced for the faithfulness comparisons.Source: §C.6 · pp. 49–51
  • VLM position bias is 8.7%, below the 15% rule of thumb, and the judge never answered T.Human participants chose A in 48–54% of sided answers; excluding 17 highly repetitive respondents changed no cut by more than 2.4 points and no share at all.Source: §C.6 · pp. 50–51

C.7 WALL-CLOCK MEASUREMENT

  • Wall-clock measurements use one H100 and remove measurement overhead; ToMe, FeatSim, and RTI are nondeterministic, while PixelDiT is timed without CUDA graphs.Two runs of the baseline methods differ in image output, and PixelDiT does not replay correctly as a CUDA graph.Source: §C.7 · p. 52, Table 21 · p. 52

D ABLATIONS ON JIT-B

  • The ablations vary one factor at a time on JiT-B and report FID-10k, which is not comparable with the main tables’ FID-50k.Development reference runs use JiT-B/32 with a uniform schedule, 50 steps, and guidance 3.0; JiT-B/16 has only a base-model result.Source: D ABLATIONS ON JIT-B · p. 53, §D.1 · p. 53, Table 22 · p. 53

D.2 A PARALLEL PIXEL-STREAM DECODER (DEVELOPMENT JIT-B/32)

  • A development ablation tracks a parallel pixel-stream decoder across training lengths, with Table 23 reporting FID-10k.The supplied passages do not state the table’s values.Source: §D.2 · p. 53, Table 23 · p. 53

D.4 GUIDANCE SCALE ALONG THE BUDGET (DEVELOPMENT JIT-B/32)

  • The development sampler ablations sweep timestep schedules and guidance along the budget, with Tables 24 and 25 reporting FID-10k.The guidance sweep runs from 3.0 to 6.0; its final row is the default uniform schedule with guidance 3.0.Source: §D.4 · p. 53, Table 24 · p. 53, Table 25 · p. 54

D.5 UPPER BOUND ON A BETTER DECODER (DEVELOPMENT JIT-B/32)

  • Decoder tests compare oracle-fine output, which supplies ground-truth fine detail in coarse regions, with downsample-lanczos output, which removes finer detail.The comparison is an upper bound on a better decoder versus resampling every output to 256^2 and back.Source: §D.5 · p. 54, Table 26 · p. 54

D.6 LAYOUT AND SAMPLER SETTINGS (FINAL JIT-B/32)

These final JiT-B/32 ablations vary layout, sampler, entropy thresholds, sampling-time layout updates, and whether token budgets are fixed per image or pooled across a batch.

D.6 LAYOUT AND SAMPLER SETTINGS (FINAL JIT-B/32)

  • Final JiT-B/32 ablations vary layout and sampler settings from a 100-epoch, 512^2 default using 50 Heun steps, guidance 3.0, five dense warm-up steps, and layout updates every five steps.The layout comparison includes random placement at matched token count and an oracle layout derived from a dense draft.Source: §D.6 · p. 54, Table 27 · p. 54

D.7 SENSITIVITY TO THE ENTROPY THRESHOLDS (FINAL JIT-B/32)

  • The sampler sweep varies guidance and timestep schedules, including time shift, Karras, cosine, and two-phase schedules; the threshold sweep shifts τ2p and τ4p while bisecting δ to keep counts near 128 tokens.The threshold values τ2p = 6.0 and τ4p = 4.0 bits come from APT, not tuning.Source: §D.6 · p. 54, §D.7 · p. 55, Table 28 · p. 55

D.8 LAYOUT EVOLUTION DURING SAMPLING (FINAL JIT-B/32)

  • Layout-evolution tests use eight ImageNet classes at 128 and 64 tokens, comparing no warm-up with the default five dense steps before rebuilding layouts during sampling.Figure 41 shows one sample’s estimate and layout and the share of image whose token size differs from the final layout.Source: §D.8 · p. 55, Figure 41 · p. 55

D.9 POOLED BUDGET ACROSS A BATCH (FINAL JIT-B/32)

  • The pooled-budget arm sets one offset per batch, giving images between 16 and 253 tokens at a mean of 128; the fixed arm gives every image the same budget.Table 30 reports FID-10k and paired PSNR and LPIPS against dense images; Figure 42 shows samples ordered by allocated token count.Source: §D.9 · pp. 55–56, Table 30 · p. 56, Figure 42 · p. 56
Loading 2610.12307v1…