Elastic Token Compression for Pixel-Space Diffusion Transformers

• Natural images concentrate detail in a small part of the frame, but diffusion transformers spend the same patch-token cost at every layer and timestep, especially wasteful in pixel space.Source: §1 · p. 1
• Useful reductions must preserve contiguous geometry, content-dependent region size, position, and summarized content; similarity merging, latent bottlenecks, and token skipping each forfeit at least one.Source: §1 · pp. 1–2
• Redundancy concentrates spatially—the most-detailed 15% of patches hold half of image detail—grows as images form, and peaks in middle blocks, where pooling outperforms skipping.Source: §1 · p. 1, Figure 2 · p. 2
• A Hilbert curve makes grouping two-dimensional patches a one-dimensional cut: each contiguous run is a connected, positionable region, cut at feature changes and recut as the image forms.Source: §1 · p. 2
• RTI keeps every patch in the first and last blocks and uses learned pooling and scattering for middle-block regions, with random training budgets enabling one checkpoint to run at any budget.Source: §1 · p. 2
• At twice the throughput, RTI matches the dense baseline and is ahead of prior reduction methods.Source: §1 · p. 2

Cheaper diffusion transformers

• Most cheaper diffusion-transformer methods reduce sampling steps, leaving per-step token cost untouched, which this work targets (Peebles and Xie 2023; Ma et al. 2024; Salimans and Ho 2022).Distillation, consistency models, and caching across steps are examples of the sampling-step approach (Song et al. 2023; Luo et al. 2023; Ma, Fang, and Wang 2024).Source: §2 · p. 2

Token reduction and adaptive tokenization

• Token merging and pruning group by feature similarity or importance, while input-adaptive tokenizers require training from scratch; RTI instead partitions spatially on a pretrained checkpoint.DC-DiT is also upcyclable, but uses router-learned chunking, targets an average rather than exact budget, and operates in VAE latents; Hilbert ordering is known, but using it to define groups is distinct (Haridas et al. 2026; Wu et al. 2024).Source: §2 · p. 2

Latent bottlenecks and skipping

• RTI replaces fixed-stride pooling in hierarchical transformers with content-adaptive partitioning, and spatially anchors its READ/WRITE latents unlike global-query bottlenecks.Region-local reads cost O(N) rather than O(NR), and each region token carries a position used by frozen rotary attention; unlike Mixture-of-Depths and per-token routing, RTI summarizes rather than skips tokens (Nawrot et al. 2022; Geiping et al. 2025; Haji-Ali et al. 2026; Raposo et al. 2024; Jain et al. 2024).Source: §2 · p. 2

3.1 Preliminaries

• The 32×32 patch grid is ordered by a Hilbert curve, so contiguous runs form connected regions and can express any quadtree tiling as well as tilings no quadtree can reach.The base model has 1,024 image patch tokens at 512² resolution with patch size 16.Source: §3.1 · p. 3

3.2 Compressing Only the Core

• RTI compresses only middle core blocks: learned READ maps N patches to R≪N tokens and WRITE restores per-patch updates, while the prelude, coda, and full-length text stream remain uncompressed.The core is fit to the shorter sequence with LoRA; the prelude and coda resist compression because they respectively form features from raw patches and produce per-pixel output.Source: §3.2 · p. 3

3.3 Content-Adaptive Region Partition

• A parameter-free O(N) pass cuts the Hilbert sequence at the largest feature jumps in the current state, merging homogeneous runs and splitting detailed areas to single patches; training samples R to support any budget.Because the backbone predicts the clean image, its features track image structure; fine regions gather on objects as they resolve during sampling.Source: §3.3 · p. 3
• RTI leads all reductions on GenEval at matched budgets; B/16 at R=512 beats dense quality and throughput, while L/16 stays within 1.8 GenEval at quarter tokens (2.59×).The table also reports that L/16 holds dense quality at half its tokens.Source: Table 1 · p. 4

3.4 The Read / Write Interface

• READ forms each region token as a learned weighted average with a size embedding and averaged rotary position, using region-restricted single-query attention to emphasize informative patches.A logarithmic size index distinguishes region sizes, while averaging rotary embeddings makes position encoding coarser for larger regions.Source: §3.4 · p. 4
• WRITE scatters the core’s region-token change back as patch-specific updates conditioned on each patch’s pre-pool features, preserving a skip around the core so the coda can produce per-pixel output.READ and WRITE act only within regions and cost O(N), regardless of R.Source: §3.4 · pp. 4–5

4.1 Setup

• Experiments freeze MiniT2I-B/16 and train only identity-initialized LoRA/READ-WRITE adapters for 40K steps, sampling R from {64, 128, 256, 512}; evaluation uses four metrics and 100-step sampling.The reported metrics are GenEval accuracy and CLIPScore, PickScore, and ImageReward; speed is median full-loop images per second over three interleaved passes on an RTX 4090.Source: §4.1 · p. 5

4.2 Main Results

• At R=512, RTI exceeds B/16 dense quality and speed (87.8 vs 87.2 GenEval; 1.84×), while L/16 reaches 87.5 GenEval within dense noise at 2.11×.On B/16, RTI scores highest overall and matches or leads on five of six GenEval categories. At R=256, it runs at 2.15× throughput for 84.7 GenEval against dense 87.2; dense remains ahead on every preference score.Source: §4.2 · pp. 5–6
• At R=256, Token Skip [Raposo et al. 2024] scores 82.3 GenEval at 1.19×, while merging [Bolya et al. 2023; Bolya and Hoffman 2023] trails RTI on every preference score on both backbones.Against RTI at matched budgets, Token Skip costs 2.4 GenEval points and merging costs 4.2; merged patches average 5.2 grid cells from their group centre versus 0.5 for RTI.Source: §4.2 · p. 6
• Latent arrays [Jaegle et al. 2022; Jabri, Fleet, and Chen 2023] score 45.7 GenEval, as their frame-wide reads discard spatial position and leave content unrepresented.Their read attention spans 11.4 of a possible 12.2 grid cells. At R=256, RTI preserves layout and detail while merging loses lettering and structure; removing position costs 39.0 GenEval points versus RTI.Source: §4.2 · pp. 6–7
• A region-token step is 2.6× cheaper, giving 16 steps versus dense 6 at 1.8 seconds and a 16.8-point GenEval advantage; above roughly five seconds, more steps favor dense.The advantage persists at 2.8 seconds, while the 100-step setting matches dense quality but is not the cheapest route to it.Source: Figure 7 · p. 6, §4.2 · p. 7

4.3 Elasticity and Ablations

• The checkpoint follows its trained-budget curve at unseen R=96, 192, and 384; adaptive cuts beat evenly spaced cuts by 2.7 GenEval points at R=128, while mean/broadcast READ/WRITE costs 7.2.Replacing anchored region tokens with a latent array drops quality to the Latent Array baseline; this swap removes grouping and position together, so its cost is an upper bound on anchoring alone.Source: Figure 8 · p. 6, §4.3 · p. 7

4.4 Where the Partition Spends Its Tokens

• As sampling forms the image, fine regions align with detail and their detail ratio to coarse regions rises to roughly six, while the budget sets region-size range and image features place fine regions.The partition reads features entering the core, not the detail measure: at R=512 size entropy is low, at R=64 it is high, and it drifts toward a varied middle during sampling.Source: §4.4 · p. 7, Figure 9 · p. 7

5 Conclusion and Limitations

• Comparisons in Sec. 4.2 confirm the three predictions derived from probing redundancy and defining four requirements for token reduction.The probing addressed whether tokens are redundant, where redundancy occurs, and what a reduction must preserve.Source: §5 · p. 7
• The Hilbert curve meets all four requirements by making image regions, rather than patches, the token unit in a model trained on patches.The interface is elastic: one checkpoint serves every budget, with the operating point selected at deployment rather than training.Source: §5 · p. 7
• Speedup is capped because the prelude, coda, and text stream do not shrink; the core cannot distinguish patches within a region.Not every grouping is a Hilbert interval, and the user chooses the budget rather than the image content.Source: §5 · p. 7

A Technical Details

• Table 2 lists training and sampling settings, while the implementation rebuilds image regions at every denoising step and samples budget R at each training step.Models adapt frozen released checkpoints; training updates READ/WRITE and rank-32 LoRA adapters on trunk attention and MLP projections, while the text stream remains full length.Source: Table 2 · p. 10, A Technical Details · p. 10
• Less than 2% of compute goes to partitioning and READ/WRITE; measured time is below analytic FLOPs at every budget on B/16.Figure 10 compares one-forward-pass cost with dense forwarding on B/16; measured throughput is reported instead of FLOPs.Source: Figure 10 · p. 10, A Technical Details · p. 10

Technical Appendix

• The appendix collects additional experiments on MJHQ-30K FID, JiT class-conditional generation, fewer sampling steps, partition order, and a longer core.It also describes protocols for redundancy, grouping, and routing analyses and provides further qualitative results.Source: Technical Appendix · p. 10

A Technical Details

• All baselines share the retrofit, backbone, LoRA setup, recipe, data, and budget protocol where applicable; only their token-reduction rules differ.Feat Sim uses cosine-similarity groups, Latent Array uses learned cross-attention latents without rotary position, and Token Skip routes 12.5% of image tokens to sparse blocks (ToMe, Bolya et al. 2023; Perceiver-IO, Jaegle et al. 2022; RIN, Jabri, Fleet, and Chen 2023; Mixture-of-Depths, Raposo et al. 2024).Source: A Technical Details · p. 10

B.1 FID on MJHQ-30K

• The evaluation measures distribution-level FID on MJHQ-30K for L/16 using 30,000 prompts and one image per prompt under the paper’s Sec. 4.1 protocol.FID is computed with clean-fid against official statistics; GenEval and preference models score images individually.Source: §B.1 · pp. 10–11
• 8.6 FID at R=512 and 13.3 FID at R=256 separate RTI from Feat Sim at matched budget and throughput; RTI at R=256 also beats Feat Sim at R=512.The R=256 RTI result uses half the tokens and achieves higher speedup than Feat Sim at R=512.Source: Figure 11 · p. 11, §B.1 · p. 11

B Additional Results

• Across Figure 11’s three panels, RTI leads Feat Sim at every budget; in the MJHQ-30K panel, both have the same speed at a given budget.That panel measures speedup by wall-clock time and uses red arrows to show the grouping rule’s cost alone.Source: Figure 11 · p. 11, §B.1 · p. 11

B.1 FID on MJHQ-30K

• +2.4 FID is RTI’s gap from dense at R=512, despite GenEval remaining within noise.This is the setting with half the patch tokens.Source: §B.1 · p. 11
• The FID results indicate that compression costs distributional fidelity before compositional accuracy.The passage relates this pattern to class-conditional results below.Source: §B.1 · p. 11
• Guidance 5.0 inflates every method’s FID equally, so comparisons hold but absolute values should not be compared with guidance-tuned numbers.Source: §B.1 · p. 11

B.2 Class-Conditional Generation

• The class-conditional evaluation transfers the frozen-backbone, LoRA, Hilbert-partition, random-budget recipe to pretrained pixel-space JiT-B/16 and JiT-L/16.It follows the JiT protocol: 50K samples, guidance 3.0 on [0.1, 1.0], and 50 Heun steps; evaluation uses R∈{64, 130, 193}.Source: §B.2 · p. 11
• RTI generalizes smoothly to an untrained R=130 budget: the R=128-trained checkpoint scores 6.34 FID versus 6.23 at R=130.The checkpoint trains with R∈{32, 64, 128, 192}, while evaluation also probes R=130 and 193.Source: §B.2 · pp. 11–12
• RTI beats Feat Sim at every budget on both backbones in FID and Inception Score, including 6.23 vs 15.99 FID at R=130 on B/16.At R=64, RTI also reaches 16.69 FID using 62% of dense token compute, versus Token Skip’s 19.11 FID at 71%; at R=130 it scores 6.23 FID versus 7.02 for the depth-augmented variant at three quarters of that variant’s token compute.Source: §B.2 · p. 12
• At R=193, RTI reaches 2.42 FID against dense JiT-L/16’s published 2.36 while using 88% of dense token compute.Its higher Inception Scores at R=193 and R=130 are not by themselves evidence of better images, since that metric rewards confident, prototypical class predictions.Source: §B.2 · p. 12
• JiT-B/16’s class tokens enter at block 4, so RTI’s core starts there and leaves patch-only encoder blocks at full resolution.The redundancy probe finds about 70% of representation retained at 25% of tokens; the shallow 256² profile leaves core placement effectively free, making the block-4 choice conservative rather than profile-forced.Source: §B.2 · p. 12

B.3 Fewer Sampling Steps

• At equal cost, RTI leads below about three seconds per image, scoring 82.7 vs 65.9 GenEval at 1.8 seconds, but dense sampling leads above five seconds.At 11.2 seconds, dense scores 87.6 vs 86.3 GenEval; RTI is 2.1× cheaper at equal quality. Six-step Euler is a weak sampler, and a higher-order solver could narrow the low-cost gap.Source: §B.3 · p. 12

B.4 Ablation on Partition Order

• Hilbert runs stay connected, unlike raster runs that tear at row ends and Morton runs that jump at dyadic-block boundaries; retraining changed only the order.The comparison starts from the MiniT2I-B/16 recipe and runs for 40K steps.Source: §B.4 · p. 12, Figure 14 · p. 13
• At R=256, mean regions hold four patches and order differences lie within noise; as budgets shrink to 8- and 16-patch runs, locality increasingly determines the ranking.The ranking becomes monotone in order locality at leaner budgets.Source: §B.4 · p. 12
• At R=64, raster trails Hilbert by −0.268 PickScore (CI [−0.296, −0.239]) and −0.154 ImageReward (CI [−0.182, −0.125]); Morton is intermediate.Every reported interval excludes zero.Source: §B.4 · pp. 12–13, Figure 14 · p. 13
• Raster loses 3.2 and 4.6 GenEval points at R=128 and 64, while Morton stays at reference; raster’s row-wrap tears harm composition, whereas Morton’s jumps cost perceptual quality.The raster losses fall outside the ±1.6 run-to-run spread.Source: Figure 14 · p. 13, §B.4 · p. 13

B.5 Lengthening the Core

• At matched budget, the [1, 15] core is worse on every metric than [3, 13], but costs 0.66× the image tokens and 0.44× on attention at R=256; at matched compute, they are indistinguishable.The image-token cost corresponds to an iso-cost budget of R≈461–572.Source: §B.5 · p. 13
• The [3, 13] core is an operating point, not an optimum; the iso-cost window excludes the full-length text stream, and [1, 15] was not timed.Including text lowers the true iso-throughput budget below R=461; with only two full-resolution blocks, [1, 15] has a speedup ceiling above the 2.5× saturation point.Source: §B.5 · p. 13

C Token Redundancy and Region Analysis

• The analyses use held-out synthetic prompts across ten domains, with pretrained MiniT2I-B/16 for the first two and a trained elastic checkpoint for partition analysis.The pretrained analyses use 512×512 generation in bf16, 30 Euler steps, guidance 5.0, and averages over prompts and sampled trajectory steps.Source: C Token Redundancy and Region Analysis · p. 13
• At R=64, Hilbert preserves the face, wings, mesh, and dancer better than raster, while the orderings are close at R=128.Figure 16 also presents matched-count, matched-throughput failure cases for the frozen backbone, elastic RTI, and Feat Sim at R=128 and R=64.Source: Figure 15 · p. 14, Figure 16 · p. 14
• The analyses measure image-token states from every block or block s0=3 on a 32×32 grid, comparing Hilbert-run groups with raster-order groups and their positions.The 1,024 image tokens are the only tokens measured; block s0=3 supplies the features used to build partitions at inference.Source: C Token Redundancy and Region Analysis · p. 14

C.1 Measuring the Redundancy

• At R=256 of 1,024, the study compares adaptive cuts at the largest Hilbert-feature gaps, evenly spaced Hilbert cuts, and skip selection of the R highest-scoring tokens.Adaptive and fixed replace tokens with their group mean; skip is not a grouping and replaces discarded tokens by the mean.Source: §C.1 · p. 14
• At R=256, pooling leaves 26% of the representation unrecovered, motivating RTI’s full-resolution coda and learned WRITE.EV measures linear headroom on one block’s activations, not output quality.Source: §C.1 · p. 14
• Figure 2 measures image detail concentration by sorting patches by finished-image gradient energy, and plots EV across sampling steps and blocks.The image samples reflect what the model draws; full-resolution ends are shaded, and EV indicates linear headroom in a block’s activations rather than output quality.Source: §C.1 · p. 14

C.2 Comparing the Groupings

• On shared frozen features, ours cuts the Hilbert order at its largest gaps, fixed spacing cuts evenly, and Feat Sim assigns tokens to nearest of R evenly spaced anchors by cosine similarity.Source: §C.2 · p. 14
• Latent Array’s learned softmax reads count as soft groups, whereas Token Skip forms no groups and supports only the third measurement.The latent groups use learned read queries from the trained checkpoint.Source: §C.2 · p. 15
• Contiguity is averaged spread around a group’s weighted centre; a single patch scores 0 and a uniform grid draw scores 12.2, also answering anchoring.Anchoring reflects whether a group has a well-defined position.Source: §C.2 · p. 15
• Fineness is −log2 group size, using participation ratio for soft groups; its correlation with detail is measured across groups mid-trajectory and is near zero at the noisiest steps.A fixed grid scores exactly 0 for this correlation.Source: §C.2 · p. 15
• 0.73 vs 0.74 for ours and 0.67 both times for fixed show agreement with Fig. 2c where the retention measurements overlap.Retention uses eq. (5) on each method’s own grouping, averaged per block across the core; it is a grouping-only lower bound, with no learned WRITE.Source: §C.2 · p. 15

C.3 Following the Trained Partition

• At R=256 mid-sampling, one-patch fine regions hold about 6× mean detail of regions sized 16+ patches, and fineness-detail correlation is +0.58.The patch-weighted entropy of the region-size distribution measures tiling variety; these are trained-partition behaviours, distinct from the frozen-feature size–detail measurement.Source: §C.3 · p. 15

D Additional Visual Results

• At R=128 and R=64, RTI retains more readable scene details than Feat Sim, while both lose the watch dial at R=128 and preserve smooth dunes at R=64.The sign remains readable under RTI at R=128 and partly readable at R=64; Feat Sim garbles or loses it, and also breaks the geometric lion into stripes and smooths the staircase into folds.Source: D Additional Visual Results · p. 15
• At the same budget, region tokens preserve both the subject and surroundings, unlike merging, latent arrays, or token skipping, which each lose different aspects of the scene.Merging flattens surroundings, position-free latent arrays cause subjects to fall apart, and token skipping thins the scene around salient objects.Source: D Additional Visual Results · p. 15
• Across R=512, 256, and 128, the additional MiniT2I-L/16 sweeps show RTI retaining structure where Feat Sim often distorts faces, objects, and composition.The sweeps cover twenty further prompts against a shared dense reference, with identical token counts and throughput for both methods at each budget. Differences include altered faces and poses, lost or deformed objects, and failures in small repeated colour structures; on pure texture, the methods can remain close.Source: D Additional Visual Results · p. 15, Figure 18 · p. 16, Figure 19 · p. 17