PaperlayerView exact revision on arXiv

Source-linked AI summary

Elastic Token Compression for Pixel-Space Diffusion Transformers

Eduard Zamfir, Christian Reisswig, Zongwei Wu, Yongqin Xian, Radu Timofte

arXiv:2608.29281v1cs.CV

TL;DR

Pixel-space diffusion transformers spend the same token budget on every patch, layer, and timestep despite spatially and temporally concentrated detail. RTI uses content-adaptive, Hilbert-ordered region tokens with elastic budgets, matching dense quality at 2.0× speed and staying close at 2.6×.

  • Problem

    Pixel-space diffusion transformers allocate one token per patch at the same cost in every layer and timestep, although image detail concentrates in small regions and resolves late in sampling.

  • Method

    RTI cuts a Hilbert ordering at feature-change boundaries, pools each contiguous run into a region token, and trains one checkpoint across randomly sampled region budgets.

  • Results

    RTI matches dense quality at 2.00× throughput and remains close at 2.6×, while leading prior reduction methods at matched budgets.

  • Takeaways & Limitations

    Region tokens provide an elastic interface in which the deployment budget is selected after training, while preserving connected, content-shaped spatial structure.

  • Takeaways & Limitations

    Speedup is capped because the prelude, coda, and text stream never shrink, and the core cannot distinguish patches within a region.

Abstract

from arXiv · show

Natural images concentrate their detail in a small fraction of the frame, yet diffusion models spend a full token on every patch, in every layer and at every timestep. The waste is largest in pixel-space models, with no autoencoder to absorb low-level redundancy first. Probing a pretrained pixel text-to-image transformer, we find its middle-block tokens redundant wherever the image is flat. The redundancy occupies connected, content-shaped regions, and exploiting it requires tokens with the same geometry. Cutting a Hilbert ordering of the patches provides them. Consecutive positions are always image neighbours, so any contiguous run is a connected region whose size and shape follow the content, and grouping in two dimensions becomes a cut in one. Existing reductions each lose part of this. Similarity merging scatters its groups, latent bottlenecks discard position, and skipping deletes what it should summarize. We cut where the model's features change most and pool each run into one region token. Our Region Token Interface (\method{}) adapts a diffusion model to these tokens, with the region count drawn at random during fine-tuning so one checkpoint serves every budget. \method{} leads prior reduction methods at matched budgets, matches dense quality at $2.0\times$ the speed, and stays close at $2.6\times$. The code and models are open-sourced at https://eduardzamfir.github.io/rti

1 Introduction

Pixel-space diffusion transformers spend uniform computation on patch tokens despite image detail being spatially concentrated and temporally evolving. RTI exploits this redundancy by partitioning patches into content-shaped regions and processing those regions in the model’s redundant middle blocks.

  • Motivation: Pixel-space diffusion transformers pay for low-level spatial redundancy in every block because they lack an autoencoder.Latent models absorb this redundancy once in a VAE, whereas pixel-space models repeatedly process it.
  • Design requirements: The reduction must preserve connectedness, content-adaptive size, position, and summarized information rather than deleting redundant patches.Similarity merging can scatter groups, latent bottlenecks discard grouping and position, and conditional computation deletes content that pooling preserves.
  • Evidence of redundancy: 15% of the most-detailed patches contain half of image detail, while the top half contains 88%.The concentration is measured in generated images from a pretrained pixel diffusion model.
  • Evidence of redundancy: Pooling N=1024 tokens into R=256 regions recovers a growing share of middle-block representations as images form, ahead of a uniform grid.The redundancy is strongest in middle blocks and decreases toward the first and last blocks; skipping recovers substantially less than pooling.
  • RTI approach: A Hilbert ordering turns two-dimensional grouping into one-dimensional cuts whose contiguous runs are connected regions with content-dependent sizes.Large feature gaps define cuts, producing long runs over flat areas and fine regions over detail while the partition is recut as the image forms.
  • RTI approach: RTI keeps full patch resolution in the first and last blocks, while learned pooling and scattering let middle blocks operate on randomly budgeted region tokens.The backbone remains frozen apart from low-rank adapters, so one checkpoint supports multiple region budgets.

2 Related Work

Prior approaches reduce diffusion-transformer cost mainly by shortening sampling or reducing tokens through similarity, importance, fixed-stride, or latent bottleneck mechanisms. RTI instead uses content-adaptive Hilbert-ordered regions inside a frozen transformer trunk.

  • Cheaper diffusion transformers: Step-reduction methods lower the number of diffusion sampling steps but leave per-step token cost unchanged.RTI targets the token cost within each sampling step.
  • Token reduction and adaptive tokenization: Token reduction methods group or prune tokens by feature similarity or importance rather than by spatial position.These methods include Token Merging and variants that vary reduction across layers and timesteps.
  • Latent bottlenecks and skipping: Hierarchical transformers pool by fixed stride, whereas RTI replaces fixed pooling with a content-adaptive partition on a frozen checkpoint.Quadtree and native-resolution tokenizers adapt input tokenization but impose different structural constraints than Hilbert intervals.
  • RTI interface: Figure 3 depicts RTI’s elastic interface: middle blocks process R ≪N region tokens, while the first and last blocks retain all patch tokens.READ pools regions and WRITE scatters updates back to patches; the budget is sampled during training.

3 Method

RTI compresses the redundant middle blocks of a pixel-space diffusion transformer by cutting a Hilbert-ordered patch sequence into content-adaptive connected regions, then reading and writing region tokens while preserving full-resolution processing at the ends. The interface is trained across region budgets and outperforms competing reductions at matched budgets.

  • Content-Adaptive Region Partition: A Hilbert ordering turns two-dimensional region grouping into one-dimensional cuts while guaranteeing connected, variable-sized regions with positional centers.Consecutive Hilbert positions are image neighbours, so long runs cover flat areas and short runs capture detail.
  • The Read / Write Interface: RTI places a learned READ/WRITE interface around a contiguous middle-block core, pooling patches into region tokens and scattering updates back to patches.Blocks before and after the core retain full-resolution patch tokens.
  • Content-Adaptive Region Partition: Feature-distance cuts produce long regions over homogeneous areas and fine regions over detailed areas using one O(N) pass, with the budget R specifying the region count.The cuts are recomputed from the current diffusion state, and R is sampled during training so one checkpoint supports multiple budgets.
  • The Read / Write Interface: READ forms each region token with within-region learned attention, region-size information, and an averaged positional embedding rather than a blind mean.The positional average becomes coarser as the region grows, matching the token to its spatial extent.
  • The Read / Write Interface: WRITE avoids identical updates by letting each patch specialize the shared region correction using its pre-pooling features.This preserves per-patch output despite the core operating only on region tokens.

4 Experiments

Experiments evaluate RTI’s quality–speed trade-offs, comparisons with alternative reductions, elasticity across unseen budgets, and how its partition allocates detail.

  • Main results: At R=256, RTI achieves 2.15× throughput and 84.7 GenEval against the dense backbone’s 87.2.
  • Main results: At R=512, RTI reaches 87.8 GenEval versus 87.2 for dense at 1.84× speed on B/16.
  • Comparisons: RTI outperforms alternative reductions at matched budgets, with gaps of 2.4, 4.2, and 39.0 GenEval points for deletion, scattering, and position removal.
  • Main results: At matched cost, RTI can use more sampling steps: at 1.8 seconds, 16 RTI steps outperform six dense steps by 16.8 GenEval points.
  • Elasticity and ablations: One checkpoint follows the same quality–speed curve at trained budgets and unseen R ∈ {96, 192, 384}, making the budget a post-training continuous knob.
  • Elasticity and ablations: Adaptive boundary placement costs 2.7 GenEval points at R=128, while fine regions increasingly concentrate detail as sampling proceeds.

5 Conclusion and Limitations

The paper identifies four requirements for token reduction and shows that RTI satisfies them with an elastic region interface. Its scope is bounded by unshrunk components, within-region indistinguishability, Hilbert-interval coverage, and user-chosen rather than image-adaptive budgets.

  • Conclusion: RTI’s comparisons confirm predictions derived from probing redundancy, including what token reductions must preserve.The requirements concern whether tokens are redundant, where redundancy occurs, and preservation of connected regions, position, content-adaptive size, and summarized information.
  • Conclusion: A Hilbert curve makes regions the token unit while one checkpoint supports every deployment budget.The operating point is selected at deployment rather than fixed during training.
  • Limitations: RTI’s speedup ceiling remains limited because the prelude, coda, and text stream never shrink.The core also cannot distinguish patches within a region, and not every grouping is a Hilbert interval.
  • Scope: The reported evaluation and implementation settings span additional experiments, baseline implementations, and compute accounting.The appendix covers MJHQ-30K FID, JiT class-conditional generation, fewer sampling steps, partition order, and longer cores.

B.1 FID on MJHQ-30K

On MJHQ-30K, RTI improves distribution-level FID over Feat Sim at matched budgets and throughput, with larger gains at smaller budgets. However, it does not match the dense model’s FID at half the patch tokens, and the absolute values use high guidance.

  • Evaluation: The MJHQ-30K evaluation uses clean-fid on 30,000 prompts with one image per prompt, while speedup is measured by wall-clock time.Feat Sim shares RTI’s budgets and speed positions, isolating the grouping rule in their comparison.
  • Results: 8.6 FID at R=512 and 13.3 at R=256 separate RTI from Feat Sim at matched budget and throughput.The gap widens as the budget shrinks; RTI at R=256 also reaches lower FID than Feat Sim at R=512.
  • Results: RTI does not match dense-model FID at half the patch tokens: its FID is +2.4 at R=512.GenEval is within noise at that budget, so distributional fidelity declines before compositional accuracy in this evaluation.
  • Caveat: FID values should not be compared directly with guidance-tuned numbers because guidance 5.0 is high for FID.The protocol inflates every method equally, preserving comparisons while affecting absolute values.

B.2 Class-Conditional Generation

On class-conditional ImageNet generation, RTI transfers across JiT backbones and budgets, outperforming Feat Sim and Token Skip while approaching dense quality on the larger model. Equal-cost sampling comparisons show token compression is preferable at low costs, whereas additional steps win at high costs.

  • Budget Generalization: At R=128, an unseen budget scores 6.34 FID versus 6.23 at R=130, indicating smooth quality across budgets.The model was trained with other budget values and evaluated on R=130 and R=193 as additional probes.
  • Class-Conditional Generation: 2.6× lower FID at R=130 on B/16 distinguishes RTI from Feat Sim, scoring 6.23 versus 15.99.RTI is better at every budget on both backbones in FID and Inception Score.
  • Class-Conditional Generation: At R=64, RTI uses 62% of dense token compute and reaches 16.69 FID versus Token Skip’s 71% and 19.11.At R=130, RTI scores 6.23 versus 7.02 for the depth-augmented Token Skip variant at three quarters of its token compute.
  • Class-Conditional Generation: At R=193 on JiT-L/16, RTI scores 2.42 FID versus the dense model’s published 2.36 while using 88% of dense token compute.The FID column carries the comparison because Inception Score can reward confident, prototypical predictions.
  • Core Placement: Across five JiT backbones, 25% of tokens retain about 70% of the representation, while the hourglass deepens with resolution and model size.This redundancy profile supports transfer of the interface but does not force the chosen core placement.
  • Sampling Trade-off: At equal cost, RTI reaches 82.7 versus 65.9 GenEval at 1.8 seconds, while dense reaches 87.6 versus RTI’s 86.3 at 11.2 seconds.Compression leads below about three seconds; above about five seconds, additional sampling steps are the better investment.

B.4 Ablation on Partition Order

Partition order matters increasingly as regions grow: Hilbert preserves connected runs, whereas raster tears at row ends and Morton jumps across dyadic blocks. At small budgets, these geometric differences produce distinct perceptual and compositional outcomes.

  • Partition Order: Hilbert runs are connected by construction, while raster and Morton runs can tear at row ends or dyadic-block boundaries.The partition always cuts each chosen order into contiguous runs, so locality determines the image geometry of regions.
  • Budget Dependence: At R=256, all three orders perform within noise because a mean region contains four patches.As budgets shrink and regions grow to eight or sixteen patches, order increasingly determines region coverage.
  • Results: At R=64, raster is 0.154 ImageReward below Hilbert, while Morton lies between them and every interval excludes zero.GenEval shows raster losses of 3.2 and 4.6 points at R=128 and R=64, while Morton stays at the reference.
  • Interpretation: Raster’s row-end tearing damages composition, whereas Morton preserves composition while costing perceptual quality relative to Hilbert.The comparison isolates the effect of run geometry at small budgets.
  • Core-Length Control: At matched compute, the longer core is indistinguishable from the shorter core on every metric despite being worse at matched budget.The longer core is cheaper, so budget-matched quality comparisons are not fair without accounting for compute.

C Token Redundancy and Region Analysis

The analyses measure redundancy in pretrained pixel-space transformer tokens and compare adaptive, fixed, and skipping reductions using retained representation variance. Visual orderings are also examined at extreme budgets under matched prompts and seeds.

  • C.1 Measuring the Redundancy: The protocol uses pretrained MiniT2I-B/16 activations, with 1024 image tokens arranged on a 32×32 grid and measured across sampling steps and blocks.The analyses use held-out prompts and capture either every block or the input to block s0=3, depending on the measurement.
  • C.1 Measuring the Redundancy: Figure 15 compares Hilbert and raster orderings at R=128 and R=64 using the same prompt and seed within each column.The orderings are close at R=128; at R=64, Hilbert better preserves the face, wings, mesh, and dancer structure.
  • C.1 Measuring the Redundancy: Figure 16 compares the frozen backbone, elastic RTI, and Feat Sim at R=128 and R=64 with identical token counts and throughput per budget.At the smallest budgets, RTI preserves more readable or separable scene elements, although both methods lose some details such as the watch dial.
  • C.1 Measuring the Redundancy: Retained variance measures how much representation remains after replacing each token with its group mean.EV=1 means exact retention, while EV=0 means the pooled tokens retain no more than a constant component.
  • C.1 Measuring the Redundancy: The comparison includes adaptive Hilbert cuts at the largest feature gaps, evenly spaced cuts on the same order, and token skipping that retains only the highest-scoring tokens.At R=256, adaptive and fixed grouping each operate on 1024 input tokens, whereas skipping is not a grouping.

C.2 Comparing the Groupings

The grouping comparison evaluates methods on common frozen features using representation retention, spatial compactness, anchoring, and content-dependent region size. Hilbert ordering retains the most core representation while producing the most compact groups.

  • C.2 Comparing the Groupings: All methods form reductions from the same frozen features, while Latent Array uses soft attention groups and Token Skip forms no groups.Ours cuts the Hilbert order at feature gaps; Feat Sim assigns tokens to evenly spaced cosine-similarity anchors.
  • C.2 Comparing the Groupings: A group’s centre is its position-weighted mean, and its spread measures the average spatial distance from grouped patches to that centre.A single patch has zero spread, while a uniform draw over the grid has spread 12.2; the same positional measure captures anchoring.
  • C.2 Comparing the Groupings: Region fineness is defined as −log2 of group size, with soft-group size measured by participation ratio, and is compared with patch detail.Fixed grids have zero size–detail correlation by construction, while the reported correlation is evaluated mid-trajectory.
  • C.2 Comparing the Groupings: Hilbert retains 0.735 of the core representation versus 0.720 for Morton and 0.697 for raster, while its groups average 1.49 grid cells versus 1.73 and 3.85.The comparison links compact runs with stronger representation retention; Morton is the informative alternative because its dyadic cells are contiguous.
  • C.2 Comparing the Groupings: In the trained partition, fine regions contain about 6× the mean detail of coarse regions and fineness correlates with detail at +0.58 for R=256.These values describe the learned partition at mid-sampling, rather than the frozen-feature size–detail analysis.

D Additional Visual Results

Additional visual comparisons test RTI against feature-similarity merging, latent arrays, and token skipping across prompts, budgets, and scene types. RTI preserves subjects and surroundings more reliably at matched low-token budgets.

  • D Additional Visual Results: At R=128 and R=64, RTI preserves more readable or separable elements than Feat Sim in several scenes, though both methods lose some details.RTI keeps parts of lettering and geometric structure, while Feat Sim garbles lettering and breaks forms; both retain smooth dunes at R=64.
  • D Additional Visual Results: At R=256, region tokens preserve both the salient object and its surroundings, whereas merging flattens surroundings, latent arrays lose subject position, and skipping thins the scene.The comparison uses three prompts and matched budgets with each method’s own trained checkpoint.
  • D Additional Visual Results: The extended sweep evaluates twenty further MiniT2I-L/16 prompts at R=512, R=256, and R=128 against a shared dense reference.The scenes include photographic, painterly, bright photographic, and saturated examples.
  • D Additional Visual Results: At R=256 and R=128, visual failures from merging include altered objects, anatomy, lettering, composition, and fine structure that RTI more often keeps recognizable.Examples include distorted meshes, glasses, riders, chairs, feeders, carousel elements, and macaw features.
  • D Additional Visual Results: The budget labels report 2.11× throughput at R=512, 2.59× at R=256, and 2.86× at R=128 for the compared visual rows.These rows compare RTI and Feat Sim against the dense baseline under the same displayed budgets.
Loading 2608.29281v1…