Decoding Register Tokens Preserves Global Scene Details
Authors: David Jacob Drexlin, Julius Hense, Jonas Dippel
Organizations: Machine Learning Group, Technische Universität Berlin · BIFOLD – Berlin Institute for the Foundations of Learning and Data Franklinstraße 28/29, 10587 Berlin, Germany · Aignostics, Germany
Pretrained vision transformers learn structured image representations spanning spatial patch tokens and non-spatial tokens. Recent analyses suggest that non-spatial register tokens encode global scene properties such as illumination, blur, and style. Yet patch-based image decoders discard these outputs, leaving the decoder to recover image appearance from patches alone. We investigate whether retaining register features improves decoding when the complete spatial patch representation is already available. We study this question using representation autoencoders (RAEs) with frozen register-equipped encoders. Our decoder processes four register tokens together with all patch tokens through joint self-attention. Appearance probes, paired reconstruction-error analysis, and fixed-decoder register substitutions examine which global image properties are accessible from registers and how their values affect reconstructed pixels. For generation, we train a multi-head attention pooling (MHAP) predictor to recover register features from sampled patches, keeping the pretrained diffusion model and its samples unchanged. On ImageNet-1K, our ViT-B decoders achieve mean reconstruction FIDs of 0.472 with DINOv2 and 0.442 with DINOv3, reductions of 9.5% and 5.2% relative to their respective patch-only baselines. The corresponding PSNRs are 19.73 and 20.54 dB. For generation, our DINOv2 models achieve a mean FID of 1.641, 9.8% below the patch-only baseline on identical sampled patches.
Figures & tables
Figure 1: Decoding patches and registers. (a) The patch-only decoder DP receives all spatial patches; DR additionally receives four encoder registers. (b) MHAP recovers registers from sampled patches before decoding with DR . Both generation routes use the same pretrained DiT and sampled patches.
Features
Thumbnail
Color hist.
DCT DC
DCT AC
Patch mean
0.211
0.288
0.735
0.038
Patch mean + registers
0.228
0.296
0.820
0.068
Table 1: ImageNet-100 appearance probes. Held-out R2 . DC measures spatial mean; AC contains nonzero low-frequency DCT coefficients.
Figure 2: Reconstruction from four registers alone. Originals (top) and reconstructions from four registers without patches (bottom). Color and texture are partly recovered, while substantial object and spatial detail is missing. Decoder and selection details appear in Appendix B.3 .
Figure 3: Appearance changes through register interpolation. (a) A blue-hour city receiver with registers from daytime towers. (b) A dog receiver with registers from a red-lobster image. Within each row, the receiver supplies the same complete patch sequence throughout. Each of the four register slots is interpolated separately toward its corresponding donor slot, using a shared interpolation coefficient α . Decoder weights and token positions remain fixed. At α=0 , decoding uses the receiver’s native registers; at α=1 , it uses the donor’s registers. Five of nine evaluated interpolation values are displayed. (c,d) Mean linear-light luminance and variance-normalized gradient energy for the two individual trajectories. Both rows use an ImageNet-1K-trained DINOv2 ViT-B decoder.
Encoder
Decoder input
rFID ↓
PSNR (dB) ↑
DINOv2
Patches
0.5211±0.0043
19.4901±0.0072
Patches + registers
0.4717±0.0032
19.7262±0.0012
DINOv3
Patches
0.4669±0.0026
20.2994±0.0182
Patches + registers
0.4424±0.0014
20.5406±0.0156
Table 2: ImageNet-1K reconstruction. Mean ± sample standard deviation across three ViT-B decoder seeds. rFID measures distributional similarity; PSNR measures pixel agreement with the source image.
Figure 4: Reconstruction and generation with register inputs. (a) Original images, patch-only reconstructions, reconstructions with native encoder registers, and signed reconstruction-error reduction. Red indicates lower squared RGB error with registers; blue indicates lower error with patches alone. (b) Patch-only generation, generation with MHAP-predicted registers, and the unsigned per-pixel RGB RMS difference between these outputs. Each generation pair uses exactly the same stored patch tensor from the pretrained DiT-DH/XL. Generation difference maps show the magnitude of change, not improvement, because no paired ground-truth image exists. Reconstruction error maps use a fixed range of [−0.04,0.04] ; generation difference maps use [0,0.20] . Examples are illustrative comparisons between separately trained decoders. Model and sample identities appear in Appendix F.6 .
Decoder
Decoding inputs
gFID ↓
ViT-B
Patches
1.8191±0.0058
ViT-B
Patches + recovered registers
1.6414±0.0133
Released ViT-XL
Patches
1.5572
Table 3: Decoding identical pretrained-XL samples. ViT-B results are mean ± sample standard deviation across three decoder seeds, with one fixed MHAP predictor for register recovery. The released ViT-XL decoder is a larger-model reference.
Appendix figures & tables21 assets
Supplementary material from the paper’s appendix.
Appendix
Batch
Route
Latency (ms/batch)
Peak memory (MiB)
1
Patch-only
9.610
377.3
1
MHAP + registers
11.577
439.3
8
Patch-only
9.420
477.5
8
MHAP + registers
11.717
541.4
Appendix
Table 4: Decoder-stage inference cost. Median CUDA-event latency per batch from 100 forwards after 20 warmups. Each route and batch size uses a fresh process; the two routes use identical stored GPU-resident patch inputs at each batch size. Peak allocated memory includes route weights, inputs, temporaries, and the live output, but excludes allocations outside PyTorch’s CUDA allocator.
Target
Patch mean
+ reg4
+Register bank
Thumbnail
0.211
0.223
0.228
Color histogram
0.288
0.316
0.296
DCT DC
0.735
0.827
0.820
DCT AC
0.038
0.040
0.068
Appendix
Table 5: Individual-register and register-bank probes. Held-out R2 from mean-pooled patches, with the selected register, or with all four registers. DC denotes the spatial mean; AC denotes nonzero low-frequency DCT coefficients.
Target
Full patches
+Registers
ΔR2
RGB DC / global mean
0.57336
0.66280
+0.08944
Very-low DCT
0.46391
0.46386
-0.00005
Low DCT
0.16085
0.16082
-0.00003
Mid DCT
0.00466
0.00437
-0.00029
High DCT
-0.00237
-0.00242
-0.00005
Thumbnail
0.24129
0.24711
+0.00582
Appendix
Table 6: Full-patch appearance probes. Held-out R2 , averaged across three shared outer splits. Changes are computed before rounding.
Figure 5: Additional reconstructions from four registers alone. Each adjacent pair shows an original photograph and its register-only reconstruction. Examples show an illuminated ship, honeycomb texture, foliage, cinema lighting, plumage, and blue sky. The ImageNet-100-trained DINOv2 ViT-B decoder receives four encoder registers, without spatial patch features or encoder CLS. All outputs use decoder seed 0 and EMA weights after 15 completed epochs.
Decoder input
DINOv2 rFID ↓
DINOv3 rFID ↓
P
4.8388
4.49223
P + CLS
4.7474
4.37634
P + REG
4.6113
4.33774
P + REG + CLS
4.5932
4.33395
Appendix
Table 7: Non-spatial decoder inputs. Reconstruction FID in the DINOv2 and DINOv3 input studies. P denotes the complete spatial patch sequence and REG denotes the four encoder registers. Comparisons are made within each encoder’s experiment.
Decoder
Epochs
Patches
+CLS+4REG
Reduction
ViT-B
16
5.7289
5.3722
6.2%
ViT-L
16
5.4250
5.1241
5.5%
ViT-XL
16
5.4720
5.1724
5.5%
ViT-L
40
4.6091
4.4234
4.0%
Appendix
Table 8: Decoder capacity with DINOv2 features on ImageNet-100. Reconstruction FID for patches alone and for patches with CLS and four registers. Epochs denote the recorded training-schedule length. The 16-epoch B/L/XL values are recorded manuscript measurements without retained seed-level metric records. The 40-epoch ViT-L row is a separately retained seed-0 evaluation. Relative reductions are computed within each row.
Approach
Parameters (r/g)
rFID ↓
gFID ↓
Patch-only
86.435 / 86.435
0.5253
1.8189
Internal InfoNCE
143.806 / 143.806
0.5104
1.7888
Internal register MSE
88.215 / 88.215
0.5181
1.8399
Native / MHAP registers
87.029 / 102.395
0.4708
1.6561
Appendix
Table 9: Register supervision and register inputs. DINOv2 ImageNet-1K, decoder seed 0, 16 epochs. Parameter counts are in millions for reconstruction/generation, excluding the frozen encoder, diffusion generator, and training-only prediction head. Counts include fixed positional parameters and, for InfoNCE, 3,072 stored null-slot parameters unused during ordinary inference.
Decoder input
rFID ↓
rFDD ↓
Patches
4.8284
41.3454
Patches + 5 learned tokens
4.7992
41.3672
Patches + 5 PatchMAP summaries
4.8168
40.9061
Appendix
Table 10: Decoder-context ablations on ImageNet-100. Single-seed measurements with five added tokens. Neither context condition receives non-spatial encoder outputs. The ablation decoders are not token-count matched to the four-register decoder.
Decoder input
rFID ↓
Patches
4.8388±0.0302
Patches + registers
4.6113±0.0087
Appendix
Table 11: ImageNet-100 decoder-input comparison. Mean and reported standard deviation across three decoder seeds. Both routes receive all spatial patch features; the second also receives the encoder’s four register outputs.
Encoder
Decoder seed
Patches rFID ↓
+Registers rFID ↓
DINOv2
0
0.525288
0.470804
1
0.516679
0.469105
2
0.521390
0.475329
Mean ± SD
0.521119±0.004311
0.471746±0.003217
DINOv3
0
0.464152
0.443066
1
0.467069
0.440756
Appendix
Table 12: ImageNet-1K reconstruction by decoder seed. All decoders use the ViT-B architecture. Register-input decoding supplies four encoder registers and the complete patch sequence. The aggregate rows reproduce the principal reconstruction comparison.
Encoder
PSNR (dB) ↑
DISTS ↓
FWD ↓
DINOv2
19.7262±0.0012
0.10444±0.00005
0.4579±0.0105
DINOv3
20.5406±0.0156
0.09670±0.00007
0.4443±0.0050
Appendix
Table 13: Fidelity of ImageNet-1K register-input reconstructions. Mean ± sample standard deviation across three decoder seeds. PSNR and DISTS compare each reconstruction with its source image; FWD compares image distributions in wavelet coordinates. Both rows use patches and four encoder registers.
Encoder
Decoder input
rFID ↓
PSNR (dB) ↑
DISTS ↓
NINCO
DINOv2
Patches
4.9567
19.9535
0.12304
Patches + registers
4.6502
20.2624
0.11652
DINOv3
Patches
4.6931
20.7924
0.11497
Patches + registers
4.4501
21.1945
0.10909
Places365
Appendix
Table 14: Reconstruction outside ImageNet-1K. ImageNet-1K-trained ViT-B decoders evaluated after 16 epochs, without training on either evaluation dataset. All results use decoder seed 0.
Evaluation dataset
rFID ↓
rFDD ↓
ImageNet-R
2.44
22.07
NINCO
6.62
65.76
Places365
4.88
26.30
Appendix
Table 15: ImageNet-100-trained patch-only OOD references. DINOv2 patch-only decoders trained on ImageNet-100. The table preserves the reported reference scores without adding unmeasured register-decoding comparisons.
Register source
Mean-color MSE ( ×10−3 )
Centered MSE ( ×10−3 )
DISTS ↓
Encoder registers
0.5665
12.7735
0.10441
MHAP predictions
1.6368
13.0146
0.11205
Training-set mean
4.5069
13.5449
0.12890
Shuffled MHAP
6.2529
13.8718
0.13331
Appendix
Table 16: Register recovery with a fixed decoder. One decoder–predictor pair, 50,000 ImageNet-1K images, and identical receiver patches. The MSE columns correspond to the two terms in Equation 5 ; lower values indicate smaller errors.
Register source
gFID ↓
Change from aligned
Aligned MHAP
1.6561
—
Within-class shuffled MHAP
1.7981
+8.57%
Training-set mean
2.3155
+39.81%
Appendix
Table 17: Register replacements during generation. One fixed decoder–predictor pair and identical sampled patches. Relative changes use the aligned seed-0 score as denominator.
Decoder seed
Patch-only ViT-B gFID ↓
MHAP + REG4 ViT-B gFID ↓
0
1.818943
1.656117
1
1.813372
1.637629
2
1.824930
1.630376
Mean ± SD
1.819082±0.005780
1.641374±0.013273
Released ViT-XL decoder
1.557247
Appendix
Table 18: DINOv2 generation by decoder seed. All rows decode the same 50,000 pretrained-XL patch samples. MHAP uses predictor seed 0, selected at epoch 96, for every register decoder. Aggregate rows report mean ± sample standard deviation across decoder seeds. The released decoder is a separate ViT-XL reference.
Generation and decoding route
gFID ↓
Patch generation, patch-only decoding
5.5981
Patch generation, MHAP auxiliary recovery
5.4041
Joint generation, λaux=0.05
7.7431
Joint generation, λaux=0.25
6.8716
Joint generation, λaux=0.50
7.1076†
Appendix
Table 19: Joint auxiliary generation. Separate DiT-S experiment with five auxiliary tokens (CLS and four registers), τ=0.8 receivers, and 50,000 evaluated images. † The retained metric row does not identify its FID reference; this value is not used for cross-route comparisons.
Contribution
Patch-only
Registers
Change
95% interval
Texture
0.085537
0.079748
-0.005789
[−0.006686,−0.004891]
Structure
0.037023
0.035759
-0.001263
[−0.001549,−0.000977]
Total DISTS
0.122559
0.115507
-0.007052
[−0.008227,−0.005877]
Appendix
Table 20: DISTS component changes on ImageNet-100. Values average the three decoder-seed means. Changes are register-input minus patch-only. Confidence intervals use paired seed differences and Student’s t with two degrees of freedom.
Statistic
Donor
n
Aligned
Monotonic
Median Δ
Luminance
Lower
10
10
7
-0.092725
Luminance
Higher
10
9
7
+0.049684
Normalized gradient energy
Lower
10
10
8
-0.203004
Normalized gradient energy
Higher
10
10
9
+0.065271
Appendix
Table 21: All retained register-interpolation trajectories. DINOv2 four-register decoder trained on ImageNet-1K, decoder seed 0. Donor indicates a lower or higher native reconstructed statistic. Aligned counts donor-directed endpoint changes; monotonic counts trajectories following that direction across all eight adjacent steps. Medians include adverse responses.
Figure 6: Additional ImageNet-1K reconstructions. Columns show the original image, patch-only reconstruction, reconstruction with four encoder registers, signed RGB MSE reduction, and RGB RMS change. Signed maps show patch-only minus register-input squared error against the original, averaged over RGB channels: red indicates lower error with registers and blue lower error with patches alone. RMS maps show unsigned differences between reconstructions. Fixed scales are [−0.04,0.04] for signed error reduction and [0,0.20] for RMS change; values outside these ranges saturate for display.
Figure 7: Additional generation examples. Each row decodes the same stored pretrained DiT-DH/XL patch tensor using a patch-only decoder or a decoder supplied with MHAP-predicted registers. The right column shows the per-pixel RGB RMS difference between outputs on a fixed [0,0.20] scale. Brighter values indicate larger output changes; values above 0.20 saturate for display. These maps measure change magnitude, not a difference in generation quality.
Attention-based models often develop attention sinks, where a small number of tokens repeatedly attract attention and accumulate unusually large activations. In vision transformers, these outliers are closely related to registers, which have been diagnostically linked to global, low-frequency image structure. Existing work has largely studied registers through interpretability analyses and linear probes, leaving open whether they can be operationalized as plug-and-play signals for generation without retraining. We revisit this question in tokenized image generation. Using OpenCLIP and DINOv2 on ImageNet, we find that test-time register features exhibit stronger low-frequency concentration than both [CLS] readouts and patch-mean features, and show a consistent (albeit moderate) correlation with pixel-space DCT low-frequency energy. Motivated by these diagnostics, we introduce RegToken, a training-free procedure that converts register structure into a small set of global prior tokens by (i) NFN-based layer localization, (ii) TokenRank-guided subspace extraction, and (iii) a projection-and-conservation update on the register subspace. Inserted into a frozen compact 1D token generation pipeline, RegToken improves ImageNet generation and alignment metrics (e.g., FID-5k 20.5 to 20.1, SigLIP 3.6 to 3.9) without modifying pretrained weights, and accelerates test-time optimization (Steps@τ 74 to 52). Overall, our results suggest that structures often viewed as attention artifacts can be repurposed as lightweight global priors for tokenized generation.
Cheng-Yao Hong, Yifan Wang, Yuewei Lin +1
Stony Brook University · Brookhaven National Laboratory
Vision Transformers (ViTs) are known to exhibit high-norm patch-token outliers that degrade feature map quality, a problem effectively mitigated by register tokens. As diffusion models increasingly adopt transformer architectures and move toward pixel-space training, they become closer in form to ViTs, raising the question of whether register tokens are also useful for Diffusion Transformers (DiTs). In this work, we show that DiTs differ from ViTs in a key respect: they do not exhibit patch-token outliers but still benefit from registers. Interestingly, registers are more effective in pixel-space DiTs than in latent-space DiTs. By analyzing intermediate representations, we find that register tokens produce cleaner feature maps at high noise levels, which may contribute to their effectiveness in pixel-space generation. We further observe that recent pixel-space DiT architectures implicitly incorporate register-like mechanisms, which may partially account for their strong empirical performance. Motivated by these observations, we propose Register Guidance, a technique that amplifies the contribution of register tokens responsible for improving visual structure and coherence.
Nikita Starodubcev, Ilia Sudakov, Ilya Drobyshevskiy +2
We study outlier tokens in Diffusion Transformers (DiTs) for image generation. Prior work has shown that Vision Transformers (ViTs) can produce a small number of high-norm tokens that attract disproportionate attention while carrying limited local information, but their role in generative models remains underexplored. We show that this phenomenon appears in both the encoder and denoiser of modern Representation Autoencoder (RAE)-DiT pipelines: pretrained ViT encoders can produce outlier representations, and DiTs themselves can develop internal outlier tokens, especially in intermediate layers. Moreover, simply masking high-norm tokens does not improve performance, indicating that the problem is not only caused by a few extreme values, but is more closely related to corrupted local patch semantics. To address this issue, we introduce Dual-Stage Registers (DSR), a register-based intervention for both components: trained registers when available, recursive test-time registers otherwise, and diffusion registers for the denoiser. Across ImageNet and large-scale text-to-image generation, these interventions consistently reduce outlier artifacts and improve generation quality. Our results highlight outlier-token control as an important ingredient in building stronger DiTs.