Persistence Forcing: Exploiting Feature Specialization in Pixel-Space Diffusion

AI summary

Authors: Chong Wang, Zixuan Fu, Shiqi Huang, Siyuan Yang, Hao Cheng, Bihan Wen (corresponding author) Affiliations: Nanyang Technological University, KTH Royal Institute of Technology, Hebei University of Technology arXiv: 2609.36014v1 [cs.CV], 28 Sep 2026

Overview

Research area: Computer Vision — pixel-space diffusion Transformers (DiTs) for class-conditional image generation.

Technical level: Advanced. The paper assumes familiarity with diffusion transformers, classifier-free guidance, and representation learning terminology.

Scope: The paper shows that assigning different feature groups different amounts of refinement across network depth in a pixel-space diffusion Transformer produces an ordered "persistent–active" feature specialization, and proposes a method (Persistence Forcing) that exploits this specialization for both denoising and sampling-time guidance.

What This Paper Is About

Pixel-space diffusion Transformers process high-dimensional raw image data directly, and in standard designs (such as JiT) every hidden feature dimension is transformed by every Transformer block, meaning all features receive the same amount of "refinement." The authors argue this uniform treatment is mismatched to natural images, where global structure can be represented compactly while fine textures require richer representation. Their goal is to let different feature groups receive different refinement budgets, then use the resulting representation organization to improve generation quality.

Key Contributions

1. Heterogeneous refinement in pixel-space diffusion Transformers. The authors let each Transformer block actively transform only the first $d_\ell \le D$ feature dimensions while the remaining dimensions bypass the block unchanged, creating different refinement histories across the hidden representation. Refinement budget per feature is defined as $R_i = \sum_{\ell=1}^{L} \mathbf{1}[i \le d_\ell]$, which equals $L$ for all features under uniform refinement.

2. Discovery of an ordered, emergent feature specialization. Sparsely refined features predominantly encode coherent global visual structure ("persistent features"), while more frequently refined features specialize toward localized, high-frequency detail ("active features"). The authors state that these roles are not explicitly assigned — the architecture determines only refinement history, and the role emerges through denoising.

3. Persistence Forcing (PerF). Persistent features are used as an explicit structural context that continuously conditions actively refined features. This is implemented as a token-wise modulation added to the standard AdaLN timestep/class modulation: $\mathbf{m}\ell = \mathrm{AdaLN}\ell(t,y) + \mathcal{P}\ell(\mathrm{Norm}(\mathbf{h}^p\ell))$, with a low-rank bottleneck.

4. Persistence Guidance (PG). Because the persistent-conditioning pathway can be stochastically disabled during training, the same model can predict with and without persistent context. Contrasting these predictions yields an internally derived guidance direction, which the authors show is complementary to classifier-free guidance (CFG).

Main Findings

• Heterogeneous refinement alone is not enough. In a parameter-matched comparison at 131M parameters, the uniform JiT variant reached FID 42.1 / IS 37.6 at 100 epochs and FID 33.1 / IS 49.8 at 200 epochs, while the heterogeneous variant reached FID 42.7 / IS 36.2 at 100 epochs and FID 33.8 / IS 47.1 at 200 epochs — i.e., it did not outperform the uniform baseline and could be worse.

• Adding the interaction helps, and guidance helps most. Progressive construction on PerF-B (100 epochs, no REPA): heterogeneous-only backbone at 131M gives FID 42.71 / IS 36.2 without CFG and 8.02 / 165.6 with CFG; adding persistent-to-active conditioning (137M) gives 40.35 / 38.2 and 7.19 / 173.8; adding Persistence Guidance (137M) gives 24.92 / 62.1 and 5.81 / 195.8.

• Consistent gains over JiT at 256×256. PerF-B reduces FID from 3.66 to 2.81, PerF-L from 2.36 to 1.91, and PerF-H from 1.86 to 1.63, corresponding to relative reductions of approximately 23%, 19%, and 12%. PerF-H also raises Inception Score from 303.4 to 324.5, and recall improves at all three scales (PerF-B 0.56 vs JiT-B 0.52; PerF-L 0.64 vs 0.60; PerF-H 0.65 vs 0.62) while precision stays comparable.

• Parameter efficiency. The abstract notes PerF-L reaches FID 1.91, approaching the 1.86 of JiT-H with only half the parameters (PerF-L is 471M versus JiT-H at 953M). PerF models introduce less than 5% additional parameters over their JiT counterparts (PerF-B 137M vs 131M; PerF-L 471M vs 459M; PerF-H 987M vs 953M).

• Gains transfer to 512×512. PerF-H/32 (992M) reduces FID from 1.94 to 1.76 versus JiT-H/32 (956M) and raises Inception Score from 309.1 to 335.3, with recall improving from 0.61 to 0.64.

• Robust to the specific width schedule. Across four contraction–expansion schedules, FID varies only from 40.35 to 41.35, compared with 42.1 for JiT. Schedules with a later bottleneck performed slightly better, with the Late and Very Late variants reaching FIDs of 40.35 and 40.44; the Late Bottleneck schedule is the default.

• Persistence Guidance is complementary to CFG, not a substitute. Across all tested CFG scales, a moderate amount of PG consistently improves FID, with the best region around $w_p = 3$. Applying PG over the full denoising interval [0,1] gave the best overall performance both with and without CFG (FID 16.03 / IS 90.3 without CFG; FID 3.12 / IS 311.0 with CFG, using CFG=3.0 and PG=2.0 when guidance is enabled).

• Larger conditioning bottleneck helps more when PG is strong. With PG disabled (PG=1.0), increasing $r$ from 128 to 256 changed FID only from 40.35 to 40.06 while doubling added parameters from 5.7M to 11.5M. With PG=4.0, FID dropped from 24.92 to 22.16 and IS rose from 62.1 to 67.2. The default is $r=128$.

• Persistent features stabilize early. The appendix analysis shows a persistent group updated only at Blocks 1 and 12 (bypassing intermediate blocks) already exhibits coherent global object structure after only the first update, and remains stable through the bypass trajectory from Blocks 2 to 10 before re-entering at Block 12. Active features, updated at every block, continue to evolve and show predominantly residual and higher-frequency structure at intermediate blocks.

Methodology in Plain English

The authors start from JiT, a simple pixel-space diffusion Transformer, and change one architectural detail: instead of every block transforming the whole hidden vector, each block only transforms the first $d_\ell$ dimensions, and the rest pass through untouched. By varying $d_\ell$ across depth with a contraction–expansion profile (narrowing then widening), some feature channels end up being processed many times while others are processed only a few times. To keep the comparison fair, they approximately match the attention and MLP parameter budget of the vanilla JiT by preserving $\sum_{\ell=1}^{L} d_\ell^2 = L d^2$.

They then examine what those differently-refined features actually contribute by decomposing the final RGB prediction into additive per-feature-group contributions through the output projection. Under this setup the contributions line up in a coarse-to-fine order; under uniform refinement, no such ordering appears.

Next they wire the two groups together. The persistent (lightly refined) features are not repeatedly transformed; instead, a low-rank projection of their normalized state is added as token-wise modulation to the active features' AdaLN conditioning, so the stable global structure guides the refinement of finer details.

Finally, they train the model so that this persistent-conditioning pathway can be randomly dropped (independently of class conditioning, each with probability 0.1). At sampling time this lets the same model produce a prediction with and without persistent context, and extrapolating along that difference gives Persistence Guidance, which can be combined with CFG in a single formula: $\mathbf{x}^{\mathrm{CFG+PG}} = \mathbf{x}^p_u + w_c(\mathbf{x}^p_c - \mathbf{x}^p_u) + (w_p-1)(\mathbf{x}^p_c - \mathbf{x}^{\varnothing}_c)$.

Implementation details: rank 128 bottleneck for PerF-B/L and 192 for PerF-H; REPA is applied on persistent features with weight 0.1 only for the main experiments, with all ablations and representation analyses run without REPA; models are trained for 600 epochs in the main results and 100 epochs for ablations.

Why This Matters

Impact on research. The paper reframes an architectural degree of freedom (nonuniform layer-wise width, previously explored for efficiency and capacity allocation in decoder-only language models) as a mechanism for inducing structured representational roles in diffusion models. It also adds a new source of internal guidance to a growing line of work on guidance signals derived from within the model itself, distinct from CFG's external class conditioning. The reported results suggest that how computation is organized across the hidden representation matters, not just how much computation is available.

Real-world applications:

• High-resolution text-to-image and class-conditional image generation, where pixel-space models avoid a separate pretrained tokenizer.
• Any deployment where image quality must be achieved under a parameter or memory budget — PerF-L reaching FID 1.91 at 471M parameters versus JiT-H's 1.86 at 953M is directly relevant to on-device or cost-constrained serving.
• Generative data augmentation and synthetic dataset creation for vision tasks, where coverage (reflected by the reported recall gains) matters alongside fidelity.
• Sampling-time quality control: because PG is a sampling-time knob applied through the full denoising interval [0,1] and combined with CFG, it can be tuned post-training without retraining.

Industry relevance. The method is a drop-in modification to an existing backbone family (JiT) with under 5% added parameters, and the authors report that the benefit holds across B, L, and H scales and across 256×256 and 512×512 resolutions — properties that favor adoption over architectures requiring a full redesign.

Future Directions

• Extending heterogeneous refinement beyond this backbone. The conclusion explicitly motivates further study of heterogeneous refinement in broader diffusion architectures and applications.
• Understanding why heterogeneous refinement alone fails. The parameter-matched heterogeneous model did not beat uniform refinement, while adding the persistent-to-active interaction did; the underlying reason for this gap is not resolved in the presented content.
• Better characterizing the emergent specialization. The depth-wise analysis in the appendix is shown for one representative persistent group and its corresponding active group at 128 dimensions; the paper's truncated content does not report a systematic study across all groups.
• Open questions on guidance design. The paper reports that PG performance degrades again when $w_p$ becomes too large, indicating the persistent context is most effective without over-constraining active refinement; the precise boundary of that trade-off, and the ranking of PG interval choices, leave room for further exploration. The provided content does not report a limitations section or compute-cost comparison beyond parameter counts.

Target Audience

Researchers and engineers working on diffusion models and generative vision, particularly those interested in pixel-space (tokenizer-free) diffusion Transformers, efficient architecture design, and sampling-time guidance. The paper is most useful to readers already comfortable with DiT-style architectures and classifier-free guidance; readers new to diffusion modeling will find the core intuition about coarse-to-fine feature roles accessible, but the equations and implementation details require background knowledge.

Authors’ abstract

Pixel-space diffusion Transformers (DiTs) directly operate on high-dimensional visual data, yet their hidden representations typically undergo uniform refinement across depth. Natural images, however, are inherently organized at different levels of granularity. Global structure can often be represented compactly, whereas local textures and fine details require richer representations. Motivated by this, we introduce heterogeneous refinement in pixel-space DiTs, assigning different feature groups distinct refinement budgets across depth. Consequently, an ordered feature specialization emerges: sparsely refined features predominantly encode global visual structure, whereas more frequently refined features increasingly specialize toward localized, high-frequency details. We refer to these two groups as persistent and active features, respectively. Building on this emergent specialization, we introduce Persistence Forcing (PerF), which explicitly exploits this persistent--active feature organization for pixel-space image generation. This enables persistent features to continuously condition actively refined features, allowing stable global information to guide the ongoing refinement of finer visual details. During generative sampling, this interaction further induces a meaningful guidance direction that promotes coherent global structure and naturally complements classifier-free guidance. On ImageNet $256\times256$, PerF-L achieves FID of $1.91$, approaching $1.86$ of JiT-H with only half the parameters, while PerF-H further achieves FID of $1.63$ and $1.76$ on ImageNet $256\times256$ and $512\times512$, respectively.