# BudgetPix: a pixel diffusion model that reads its own first draft to decide where the tokens go
> Satyajit Ghana — Head of Engineering @ Inkers Technology
> canonical: https://ai.thesatyajit.com/articles/budgetpix
> date: 2026-10-13
> tags: diffusion, diffusion-transformers, image-generation, flow-matching, tokenization, inference
AK posted the title and a Hugging Face link. A pixel diffusion transformer spends the same work on every 16×16 patch: empty sky costs what fur costs. BudgetPix asks whether one fine-tune can turn that cost into a dial you set at inference time.
The paper is **BudgetPix: Compute-Adaptive Tokenization for Pixel-Space Image Diffusion** ([arXiv 2610.12307](https://arxiv.org/abs/2610.12307)), by Ozgur Kara, Yujia Chen, Daniel Watson, David Forsyth, James Matthew Rehg, Wen-Sheng Chu and Du Tran, from UIUC and Google. It adapts released pixel-space models (JiT, MiniT2I, PixelDiT and the one-step pixel MeanFlow) and reports that one checkpoint degrades gracefully from the dense grid down to a tenth of it.
The quadtree is the easy part. The interesting part is a chicken-and-egg problem the paper states itself: to decide where an image is detailed you need the image, and a generator does not have one yet. How it gets around that decides what the method costs.
There is no code and no weights. The [project page](https://karaozgur.com/BudgetPix/) shows a "Code (soon)" button that is disabled, its tooltip reading "Code and checkpoints will be released". Hugging Face's paper page links a repository, `ozgurkara99/budgetpix`; I cloned it, and at commit `f48ce48` (12 October 2026) it holds a README and a `.gitignore`, with the line "Code will be released soon!". The Hugging Face API lists no models. Everything below is checked against the paper, its long appendix and the project page.
## Why uniform patches waste compute in pixel space
A pixel-space model skips the VAE and patchifies raw RGB, so the token count is set by one number, the patch size: a 512×512 image in 16-pixel patches is $(512/16)^2 = 1{,}024$ tokens, and in 32-pixel patches it is 256. [PixelUMM and PixelDense](/articles/pixelumm-pixeldense), from the same fortnight, explain why pixel models predict the clean patch rather than the noise; BudgetPix keeps that and changes only what a token covers.
The paper's motivating example is a flag against open sky: MiniT2I-L spends 1,024 tokens on it, and BudgetPix reconstructs the scene with 307 (reported, Figure 2). The same MiniT2I-L is the text-to-image base of [Representation Guidance](/articles/representation-guidance), from the same week, which leaves the token count alone and changes the training target instead.
How much compute does a token cut buy? Count one standard transformer block with width $D$ and $N$ tokens. The QKV and output projections cost $8ND^2$ FLOPs, an MLP of ratio 4 costs $16ND^2$, and attention scores plus mixing cost $4N^2D$. At a token fraction $f$:
$$
\frac{\text{FLOPs}(fN)}{\text{FLOPs}(N)} = \frac{6Df + Nf^2}{6D + N}.
$$
The attention term is quadratic, but big pixel patches keep $N$ small next to $6D$. With the widths in the paper's Table 16 and the token counts in Table 21, attention is about 4% of a block's FLOPs for JiT-L/32 at 512² ($N = 256$, $D = 1{,}024$), 12% for MiniT2I-L/16 at 512² ($N = 1{,}024$, $D = 1{,}248$) and 31% for PixelDiT at 1024² ($N = 4{,}096$, $D = 1{,}536$) (reasoned). So cutting tokens to 25% cuts a block's FLOPs to between 19.2% (PixelDiT) and 24.3% (JiT-L/32), not to 6.25% (reasoned). The saving is mostly linear, which matters once the paper's "compute budget" turns out to mean tokens.
## The entropy quadtree
BudgetPix borrows its layout rule from APT, an adaptive-patch vision transformer. Take token scales $p$, $2p$ and $4p$ (16, 32 and 64 pixels for a 16-pixel parent). Start from the coarsest cells. For each cell $R$, build a histogram of pixel luma and take its Shannon entropy:
$$
H(R) = -\sum_{k=1}^{K} q_k(R)\,\log_2 q_k(R),
$$
where $q_k(R)$ is the fraction of the cell's pixels in bin $k$. The appendix pins the details: $K = 512$ bins, so $H$ runs from 0 to 9 bits. A $4p$ cell stays one token if $H \le \tau_{4p} + \delta$, otherwise it splits into four $2p$ cells; each of those stays one token if $H \le \tau_{2p} + \delta$, otherwise it splits into four $p$ tokens. The thresholds are $\tau_{4p} = 4.0$ and $\tau_{2p} = 6.0$ bits, taken from APT "not tuned" (reported, App. D.7). The global offset $\delta$ is the dial.
There are two ways to set $\delta$. In threshold mode it is fixed, and each image takes whatever token count its content gives. In budget mode, which the paper uses for every reported number, $\delta$ is bisected per image so the layout has at most $B$ tokens. The two lie on nearly one quality curve when you compare them at the same mean token count (reported, Table 14). Threshold mode gives up predictability: at $\delta = +1$ on JiT-L/32 the mean is 56.0% of the grid with a standard deviation of 18.2 points across images (reported). A server with a latency target cannot take that spread.
Drag the budget down and watch which regions go first. The sky merges into 32-pixel tokens, then into 64-pixel ones; the tree crown holds out longest. At the default of 30% the scene lands on 76 tokens of 256 (measured, in the widget). In the entropy view, look at the kite. Its four panels are a handful of flat colours, so its histogram has only a few occupied bins and its cells score lower than the grass, even though the kite has the sharpest edges in the frame, and it merges before the grass does. A luma histogram counts distinct values, not their arrangement; a two-colour checkerboard scores 1 bit. The paper's conclusion lists "learning where to place tokens, not deriving it from entropy" as the next step.
Uniform patchify on this grid gives 256, 64 or 16 tokens and nothing in between. The quadtree gives any count from 16 to 256 and spends the cut where the picture has least to lose. [Level-of-Token Diffusion](/articles/lot-diffusion), on arXiv three days earlier, tiles a DiT's token grid the same way, but its layouts come from a planner or a reference image and its count is whatever the thresholds produce. BudgetPix guarantees the count and needs no outside layout; the price of that comes up below.
## Getting pixels into and out of mixed-size tokens
The central transformer is the released model, fine-tuned, with its own global position embedding evaluated at each token's centre. What is new sits on either side of it.
**The encoder** has two paths. A $p \times p$ patch goes through the parent's own patch embedder $E$, unchanged. A larger patch $P_i$ of side $np$ is bilinearly downsampled to $p \times p$ and embedded with the same $E$, which loses everything inside it. A second stream, the Patch Aggregator, puts it back: it splits $P_i$ into $n^2$ sub-patches, embeds each with $E$, adds a learned position within the patch, and mixes them with a small convolutional stack into one correction vector. Both paths add a scale embedding:
$$
\mathbf{x}_i = E\big(\mathrm{resize}_p(P_i)\big) + \mathbf{e}_{\mathrm{scale}}(s) + \mathrm{PA}_s\big(\mathrm{split}(P_i)\big), \qquad n > 1.
$$
**The decoder** mirrors it. The parent's head turns a token into a $p \times p$ patch; for a coarse token that patch is bilinearly upsampled to $np \times np$, right in layout but smooth inside. A Patch Refiner adds the missing texture as a residual: it cuts the predicted patch into 4×4-pixel pieces, runs two small self-attention blocks over them under adaptive layer norm conditioned on the token's final hidden state, and projects to an $np \times np$ residual.
The added modules are 2.39M parameters on JiT-B/16, 3.50M on MiniT2I-L's 911.8M and 4.37M on PixelDiT's 1302.5M (reported, Tables 16, 18 and 19). Their last layers are zero-initialised, so at step 0 the adapted model reproduces the released one, with identical dense-grid FID (reported, Table 17). PixelDiT already ends in a full-resolution pixel stage, so there BudgetPix corrects each cell's conditioning of that stage instead, and the stage decodes every pixel at every budget.
On JiT-B/32 at 512² and a 50% budget, FID-50k is 60.61 with neither (the released weights fed coarse tokens), 11.58 with the aggregator and 8.30 with the refiner added too (reported, Table 2). The aggregator does most of the work.
## Training one checkpoint for every budget
Every weight is fine-tuned with the parent's own $x$-prediction flow-matching loss, $\mathbb{E}\big[w(t)\lVert x - \hat{x} \rVert_2^2\big]$, with $w(t)$ the parent's weighting. The new part is which layout each training image gets:
- **dense grid**, with probability $a = 0.15$, to anchor the model;
- **random layout**, on $b = 0.10$ of the rest (8.5% of all images): the quadtree of a fake entropy map drawn uniformly from 0 to 9 bits, so coarse tokens also land on detail;
- **content layout** otherwise: the quadtree of the clean image's real entropy.
For the last two, $\delta$ is drawn from $\mathcal{U}(0, 2)$ per image and step, which the paper says covers "roughly 25% up to 100% of the original dense grid". Table 14's threshold sweep agrees: $\delta = +2$ gives a mean of 28.0% of the grid on JiT-L/32 (reported). So budgets below about a quarter ask for offsets the model never trained on (reasoned), and FID shows it: 11.54 at 25%, 18.12 at 20%, 44.62 at 10% (reported, Table 14). The introduction's "coherent, high-quality images using as little as 10%" rests on the sample grids; the number at 10% is 44.62 against 2.71 dense.
## Where the layout comes from when there is no image yet
In training, the layout is the clean image's quadtree. At sampling time there is only noise. BudgetPix's answer:
1. Run the first $n_w$ steps on the dense grid, so the model sets the composition.
2. Because the parents predict $x$ directly, every forward pass already outputs a clean-image estimate $\hat{x}$. Compute its quadtree and bisect $\delta$ to fit the budget $B$.
3. Keep that layout for $m = 5$ steps, then rebuild it from the newest $\hat{x}$.
The defaults are $n_w = 5$ for the 512² JiT parents, MiniT2I and PixelDiT, and $n_w = 10$ for JiT at 256² (reported, Tables 18, 19 and 21). The layout is not predicted from the prompt or fixed in advance. It is the model's own first draft, measured with the training rule.
The appendix tests how much each choice matters, on JiT-B/32 at 512² with FID-10k (reported, Table 27). Random placement with the same token count is far worse: 19.51 against 10.99 at 50%. An oracle, the quadtree of a finished 50-step dense draft from the same noise, is no better than the default (11.30). Rebuilding every step, every 25 steps or only once after the warm-up all land between 10.96 and 11.04. The one knob that moves FID is the warm-up length: 12.44 with none, 10.99 with five, 9.68 with twenty.
So the layout only has to be roughly right. Good, because the warm-up is not free. JiT samples with 50 Heun steps under classifier-free guidance, four network calls per step, the last step Euler: 199 calls. Five dense steps plus the estimate the first layout is built from are 21 of them, at full price (reported, Table 21).
Put the reported FIDs on a compute axis and a trade appears (reasoned, using the block FLOP ratio for $N = 256$, $D = 768$). Twenty warm-up steps at a 25% budget cost about 55% of the dense FLOPs and score 15.73. Five warm-up steps at a 50% budget cost about 54% and score 10.99. At equal compute, spend it on tokens, not on a longer warm-up.
A one-step generator has no trajectory to read an estimate from. For the one-step pixel MeanFlow model the main text says BudgetPix uses "the random layout", but the appendix describes something else: "each image takes the entropy quadtree of one of 20 training images of its class, picked by the image index and bisected to the budget" (reported, App. B.6). The one-step layouts are borrowed from real ImageNet photos of the same class: a fair prior for class-conditional ImageNet, and one that does not exist for a text prompt.
The wall-clock runs use a layout "built once, then frozen" after the warm-up, not rebuilt every five steps (reported, Table 21); Table 27 scores the two the same, so the timed sampler is close to, but not literally, the evaluated one. The paper also names the obvious failure in its conclusion: the layout comes from intermediate predictions, "so it can miss details they do not yet show".
## Checking the headline claims
### "25% of the original compute budget"
The abstract says BudgetPix matches MiniT2I-L at 512² and PixelDiT at 1024² "using just 25% of the original compute budget". Figure 1's caption settles what the percentage is: "Percentages show the fraction of tokens used." Table 21 lists the budgets as token counts. The compute budget is the token budget, applied after the dense warm-up.
For MiniT2I-L, one block at 25% of the tokens is 22.7% of the dense FLOPs (reasoned, from the formula above). MiniT2I samples with 100 Euler steps, two calls each, plus one call for the estimate: 201 calls, 11 of them dense. That puts the whole sample at about 27.0% of the dense FLOPs (reasoned; it ignores the text tokens and the added modules). The paper's own measurement is a 3.59× wall-clock speed-up on one H100, which is 27.9% of the released model's time (reported in Table 6; the inversion is mine). FLOPs and time agree here, but not everywhere. At the same 25% budget, the measured speed-ups are 1.95× for JiT-L/32, 1.95× for PixelDiT and only 1.36× for JiT-B/32 (reported, Table 6). For JiT-L/32 my FLOP estimate is about 32% of dense, warm-up included, against a measured 51% of the time (reasoned against reported). Short 256-token sequences do not keep an H100 busy on matrix FLOPs (my reading). PixelDiT's floor is its pixel stage, which the paper names as a limit.
Is BudgetPix's overhead counted? Yes, and it rounds to zero: BudgetPix at 25% on MiniT2I-L takes 1.376 s per image, exactly what the untrained base model takes at the same budget through the same sampler (reported, Table 7). Token count buys the speed. BudgetPix makes the low-token images worth having.
The cheap end of the dial buys nothing. At a 90% budget MiniT2I-L is 1.01× and JiT-B/32 0.92×, and PixelDiT is 0.79×, slower than dense (reported, Table 6). Table 21 notes that PixelDiT's attention is masked when a reduced budget pads the layout, which is my best guess at the cause. From about 80% down the gains start.
Then quality. GenEval at 25% is 0.874 against the released MiniT2I-L's 0.882, and 0.725 against PixelDiT's 0.721 (reported, Tables 1 and 5). With 553 prompts of four images each, GenEval's standard error near 0.87 is roughly 0.007 (reasoned, treating it as one proportion), so the MiniT2I gap is noise. The PixelDiT comparison flatters BudgetPix: its own dense score is 0.741, because the fine-tune on PD12M improved PixelDiT before any tokens were cut (held-out FID 43.7 to 36.2 on that data, reported in Table 17). Against its own dense grid, BudgetPix at 25% loses 0.016 GenEval and 1.5 DPG-Bench points. "Matches the fidelity" holds for PixelDiT against the released model, not against itself.
### "0.8-point increase in FID" with MeanFlow
The table is App. B.6, Table 15: the one-step pixel MeanFlow model pMF-L/16, class-conditional ImageNet at 256², 256 tokens dense. BudgetPix scores FID-50k 2.65 at 100% and 3.46 at 60%, which is 154 tokens (reported). That is the 0.81 in the introduction. The baseline is BudgetPix's own dense score; the released model reads 2.54 in the same protocol, so against it the rise is 0.92 (reasoned). And 0.81 on a base of 2.65 is a 31% rise in FID (reasoned). "Near-zero quality degradation" is generous for that.
The other rows are what make it impressive. At 60%, the released weights fed BudgetPix's layouts score 257.20, ToMe 36.04 and FeatSim 100.26 (reported). Even at 25% BudgetPix is 5.44. Nothing times the one-step model, so there is no wall-clock behind this result.
### "Human and VLM judges" against "prior budget-adaptive baselines"
The baselines are the base model (released weights pushed to the budget through BudgetPix's sampler, no training), ToMe and FeatSim (two token-merging methods wrapped around the released weights), and RTI, a LoRA adapter that consumes region tokens, run from its released checkpoint. RTI, the only real competitor, is concurrent (August 2026), not prior, and exists only for MiniT2I, so the studies run on MiniT2I-L at 512².
Both studies measure faithfulness to each method's own full-budget image, not quality. The first asks at which budget a difference first becomes visible. People (50 per form, 20 GenEval prompts) say BudgetPix can be cut by 51.8% [48.0, 55.7] before they notice, against 38.0% [31.8, 44.0] for RTI and 32.5% or less for the rest (reported, Table 20). A Gemini 3.8 Flash judge on 100 prompts gives 36.2% against 21.1% for RTI. The intervals do not overlap; that part holds.
That contest is closer. Head to head, people prefer BudgetPix to RTI 63.5%, 63.4%, 57.2% and 62.3% of the time at 80, 60, 40 and 20% budgets, ties counted as half (reported). At 40% the interval is [45.3, 69.5], which includes a coin flip. The VLM judge at 40% gives 54.5% [46.0, 63.0], also including 50%. "Significantly improved" over RTI holds at three of four budgets.
RTI does not collapse like the merging methods: at 25% on MiniT2I-L it scores 0.859 GenEval, 83.0 DPG-Bench and 1.16 ImageReward against BudgetPix's 0.874, 84.3 and 1.17 (reported, Table 5). Where RTI loses is speed, 2.16× against 3.59× at that budget, and pixel faithfulness (Figure 7, bottom row).
## What I would watch for
- **No code.** Every number here is the authors'. The appendix is unusually complete, so the claims will be checkable once code lands.
- **Entropy is not detail.** Flat graphics, text and line art have sharp structure and low entropy; the layout spends tokens on noisy texture before clean edges.
- **Tokens are not time.** With short sequences like JiT's 256, or a full-resolution stage like PixelDiT's, wall-clock gains sit well under the token cut.
- **The warm-up is a fixed tax.** About a tenth of JiT-L/32's calls run dense before the dial does anything. Fewer sampling steps shrink the trajectory the layout is read from; the one-step case already needed borrowed layouts.
- **Below 25%, you are extrapolating.** The training offsets cover roughly a quarter of the grid and up.
Neighbours on this site attack the same cost elsewhere. [LoT Diffusion](/articles/lot-diffusion) takes the layout as an input from a planner. [MrFlow](/articles/mrflow-diffusion-acceleration) moves the cost along the time axis. [EVA](/articles/eva-empirical-vae) works on the latent side BudgetPix avoids. For why token count grows with the square of the resolution in the first place, see the [Diffusion Transformer page](/architectures/diffusion-transformer#tokens-grow-with-the-square-of-the-resolution).
## How I checked
I read the paper in the arXiv HTML (v1, 8 October 2026) and, where the HTML dropped text after a less-than sign, in the project page's PDF. Figures 4, 5 and 7 are the project page's full-resolution copies, Figure 6 is from the arXiv HTML, and Figure 41 is a 300 dpi render of PDF page 55. I shallow-cloned `ozgurkara99/budgetpix` and queried the Hugging Face API.
The FLOP ratios, warm-up accounting, GenEval standard error, inverted speed-ups and MeanFlow percentages are my arithmetic on the paper's tables. The FLOP model is one standard block with an MLP ratio of 4; it ignores conditioning tokens, the added modules and kernel efficiency, so treat it as a direction. The quadtree widget implements the appendix's rule on a scene I drew; its token counts are measured on that scene only. I could not run BudgetPix or reproduce any reported number.