Vision Transformers with Self-Distilled Registers

We appreciate your detailed and concrete suggestions! We will incorporate all of your feedback into our paper.

Q1) Detailed method description

We agree that the paper could be improved with additional details about our distillation procedure. We have already incorporated your suggestions into our internal revision. Due to NeurIPS OpenReview policy this year, we cannot upload a revision.

We provide a condensed pseudocode here, where we assume we are targeting a vision transformer with patch size $k \times k$, an image $\mathcal{I}$ which is an integer multiple of the patch size. Augmentation parameters $(\theta_1, \theta_2, ..., \theta_n)$ for up to $n$ augmentations are generated randomly from a constrained distribution, and each augmentation consists of a binary horizontal flip parameter $\text{flip}_i \in \{0,1\}$, and a vertical/horizontal offset tuple $(x_i, y_i)$. We also have a tuple grid of 2D image coordinates $C$, which helps us index the locations in an image. An augmentation function takes as input $(\mathcal{I}, C, \theta_i)$ and generates the corresponding shifted/flipped image and coordinates.

Pseudo code:

Inputs: Image $\mathcal{I}\in\mathbb{R}^{H\times W \times 3}$; Image space coordinates $C\in[0,1]\times[0,1]$; Augmentation parameters $\theta_{1}, \theta_2, ..., \theta_n$; Augmentation function $T$; ViT teacher model $f_\text{teacher}$; Zero-init feature accumulator tensor $Q$; Zero-init count occurrence tensor $K$

For i in ${1,...,n}$

$\theta_i = (x_i, y_i, \text{flip}_i)$

($\mathcal{I}_i, C_i$) = $T(\mathcal{I}, C, \theta_i)$ # Augment image, and generate the image-coordinates of the shifted/flipped image

Dense Feature $F_i = f_\text{teacher}(\mathcal{I}_i)$

$(F^\text{valid}_i , C^\text{valid}_i) = T^{-1}(F_i , C_i, \theta_i)$ # Given the coordinates and the feature map of the augmented image, we map the feature back to the position in the original image

$Q[C^\text{valid}_i] = Q[C^\text{valid}_i] + F^\text{valid}_i$

$K[C^\text{valid}_i] =K[C^\text{valid}_i] + 1$

return $Q/K$

We include a proof on page 12 of the supplemental showing that this achieves the same expected value, as gradient-based MSE optimization used in DVT neural field learning.

We will include additional details in our updated revision.

Q2) Ablation study of the distillation loss

This is a good suggestion. Here we perform an ablation study using MSE, cosine similarity, and both combined (used in our work, and also used in DVT) -- in this table higher is better:

Note that cosine similarity alone does not preserve magnitude, and occasionally yields training instability when used alone.

We will include this result in the revision.

Q3) Computational efficiency

We agree that additional quantification of training and inference costs is beneficial. Given that custom kernels can significantly influence both performance and memory usage, our evaluation focuses on model size, FLOPs, and measured wall-clock training time. To ensure a consistent baseline, training time was measured using the official DVT repository code without modifications. Since the original DVT implementation specifies a single GPU for its feature extraction and learning phases, we similarly constrained all stages of our model to a single GPU for this comparison. The evaluated models (DVT and PH-Reg) are based on the same architecture: DINOv2 ViT-B.

Training wall-clock time:

During training, DVT must store serialized Instant-NGP features, requiring approximately 1.4TB of storage. In contrast, our method imposes no supplementary storage requirements.

To assess testing costs fairly, all models were evaluated with parameters cast to fp32 and using an eager attention implementation. We account for positional embeddings adjusted for 448x448 resolution into the PH-Reg parameter count. Regarding MaskCLIP and NACLIP, their official implementations currently retain the same parameter count as the original CLIP model, though MaskCLIP can achieve a minor reduction by removing the final query and key head.

In the case of DVT, we assessed their second-stage model, which pairs a transformer block denoiser with an original vision transformer.

We achieve notable efficiency gains over DVT: a 2x reduction in training time, complete elimination of training storage costs (0 TB vs. 1.4 TB), and ~10% savings in both inference FLOPs and model parameters. These metrics will be added in the upcoming revision of our paper.

Q4) Comparison against the official DINOv2 + registers

Here we perform a comparison against DINOv2 + registers, using the official weights provided by facebook research:

We find that our method achieves consistent improvements all on metrics compared to the original DINOv2 (without registers). Our method has better performance on VOC21 segmentation and NYUv2 depth prediction compared to DINOv2 + registers, while having slightly inferior performance on ADE20k. Note that the VOC21 mAcc degradation of DINOv2 + registers compared to DINOv2 was also observed by DVT. For VOC21 mIoU, DVT observes a very small improvement, while we observe a very small degradation. For fairness, we ran our own evaluation using the same resolution for all methods -- and the evaluation settings are slightly different than those used DVT in their paper.

Given that our method can be completed using around 1.5 days on 4 Nvidia Ada GPUs, we believe our framework offers a very attractive alternative to full model retraining (estimated at several GPU-months if referencing OpenCLIP ViT-B numbers). We will include this comparison in the revised paper.

Q5) Ablation studies of distillation approach

Here we evaluate the contribution of each component in OpenAI CLIP + PH-Reg (mIoU, higher is better):

By comparing our full pipeline with the Distill from teacher w/o registers w/o TTA baseline, we find that approximately half the improvement comes from registers, and half comes from test-time augmentation of the teacher network.

We will include these results in the revision.

Conclusion

We are very grateful for the highly detailed and actionable suggestions. We have included additional details for our method, additional experiments on distillation loss, method efficiency, comparison against models with registers trained from scratch, and ablation studies of all components of our pipeline.

Given our experiments and clarifications, we hope we have resolved your questions! We look forward to any additional suggestions you may have.

We are deeply grateful for your helpful comments! We will address specific questions below, and will include additional details in a revision.

Q1) Dynamic determination of artifact tokens

This is a very interesting idea. Our current framework only removes the artifact tokens, but identifying the artifact tokens dynamically could be a promising direction. In principle, a "corrector" network could be learned by estimating the difference between the original "raw" dense features and denoised teacher features. This network could be conditioned on the original dense features and learn the artifacts.

After identification, the artifacts could be either removed or shifted to a register using ideas from recent work.

We will include a discussion on this in the upcoming revision.

Q2) Magnitude of Artifact Tokens

We were not previously aware of Vision Transformers Don't Need Trained Registers. After a careful reading, we agree it is highly relevant and will include a reference to this work in our revision. It is striking that artifact tokens could be removed using heuristics alone.

We would like to clarify that for CLIP models, we focus on the zero-shot open-vocabulary segmentation task. In this case, using the value features (or some modified convex combination of the value features across the image) of the last attention block are the preferred approach [1, 2, 3, 4].

As we show in Figure 6, we find when value features (or convex combination of value features) from the last layer are examined for CLIP models, indeed the artifact tokens are primarily low magnitude, rather than high magnitude in OpenAI's CLIP and OpenCLIP. However in DFN-CLIP, which shares the same architecture and training objective as the OpenCLIP model we evaluate -- the artifact magnitudes are more mixed. The primary difference between DFN-CLIP and OpenCLIP are the training images. This suggests that training data may play an important role in the emergence of artifact tokens.

Discussions with some DINOv2 authors offline, also suggests data may be a factor. When discussing why artifacts appear in DINOv2 but not DINOv1, one DINOv2 author suggests that increase proportion of images with white backgrounds (collected from online retail) in the DINOv2 training set may be the cause, whereas DINOv1 was trained with a higher ratio of naturalistic photographs. While this is plausible, we do not believe it is a complete explanation.

We will include additional discussion and a citation to the Don't Need Trained Registers work in our paper revision.

[1] Extract Free Dense Labels from CLIP (ECCV 2022)

[2] SCLIP: Rethinking Self-Attention for Dense Vision-Language Inference (ECCV 2024)

[3] ClearCLIP: Decomposing CLIP Representations for Dense Vision-Language Inference (ECCV 2024)

[4] Pay Attention to Your Neighbours: Training-Free Open-Vocabulary Semantic Segmentation (WACV 2025)

Q3) Global information in artifact and register tokens

It is fascinating that there seems to be a diverse set of views on the nature of artifacts -- [5] argues they include global information, while [6] argues they are image independent, with very recent work by Lappe and Giese [7] finding that some dimensions of artifact tokens contain global information -- with results varying due to model size and layernorm.

Similarly, there are disagreements about what causes artifact tokens to occur. DVT suggests that the positional embeddings are primarily responsible for artifacts. However in section 4.3, we find that artifacts cannot be fully resolved by finetuning the positional embeddings alone.

We have performed a visualization of the register tokens learned using PH-Reg -- which we will include in an upcoming revision. We find that in our current setup, the attention maps for the registers do indeed correspond to semantic objects within the image. However qualitatively the effect seems weaker than the models trained from scratch. One potential reason is that we do not constrain the [CLS] token of the student to decode global image information, and instead focus on the quality of dense representations.

We will include this attention map visualization in the upcoming revision, and discuss why our results may differ from models trained with [CLS] token constraints.

[5] Vision Transformers Need Registers (ICLR 2024)

[6] SINDER: Repairing the Singular Defects of DINOv2 (ECCV 2024)

[7] Register and CLS tokens yield a decoupling of local and global features in large ViTs (Arxiv)

Conclusion

We thank you for your comments, and we hope that this clarifies our results! We will update the paper to reflect your suggestions.

Thank you for your detailed suggestions! We will incorporate all of your feedback into our paper.

Q1) Investigation of register counts across models

This is an interesting question. In the original paper we evaluate the number of registers only on the OpenAI CLIP ViT model. Based on your suggestion, we further evaluate OpenCLIP, DFN-CLIP, and DINOv2 -- three models that were trained using different datasets (in the case of OpenCLIP and DFN-CLIP) and different objectives (DINOv2). In all cases we observe decreased rates of performance gain when additional registers are added, across median cosine similarity (higher is better):

We are not the first to observe this, as similar trends are observed in the saturation of dense task performance in the original "Vision Transformers Need Registers" [1] paper. Both our own visualizations (will be added to upcoming revision), and those in [1] suggest that registers learn to focus on semantic objects within the image. We speculate that the saturation in performance is due to the limited number of objects in common images.

One potential follow-up experiment could be based on synthetic datasets like CLEVR, where the number of objects within each image is known. These datasets could be used to investigate if the performance benefit of register tokens correlates with the average number objects across a dataset.

We will include the table for different models in an upcoming revision of our paper.

Q2) Efficiency of PH-Reg compared to DVT

We agree that additional quantification regarding both training and inference cost would be useful. Due to the potential influence of custom kernels on both performance and model memory usage, here we evaluate model size, FLOPs, and wall clock training time. When evaluating training time, we utilize the official code provided in the DVT repository without any modification. The original DVT code specifies a single GPU for feature extraction and learning -- for fairness we also limit all stages of our own model to a single GPU here. For DVT and PH-Reg, we evaluate the same model (DINOv2 ViT-B) at the same image resolution.

Training wall-clock time:

During training, DVT also requires storage of the serialized Instant-NGP features. Which requires approximately 1.4TB, while our method requires no additional storage cost during training.

When evaluating testing cost for all models, we cast parameters to fp32 dtype, and use eager attention implementation for all models. For our method, we include the positional embeddings adapted for 448 resolution. For MaskCLIP and NACLIP, current official implementations have the same number of parameters as the original CLIP, although a small reduction can be achieve in MaskCLIP by discarding the last q,k head.

For DVT, we evaluate their stage 2 model, corresponding to the transformer block denoiser coupled to a original vision transformer.

Our method is significantly faster to train (2x speed), and utilizes significantly less storage during training compared to DVT (0 vs 1.4TB). Our method also utilizes approximately 10% fewer FLOPs and 10% fewer parameters during inference compared to DVT. We will include these details in an upcoming revision of our paper.

Q3) Why so well on Vision Language Models?

Indeed, our performance gains are the most striking on models trained with an image-language contrastive loss, where OpenAI CLIP, OpenCLIP, and DFN-CLIP were trained on different datasets. While we observe performance improvements on DINOv2 compared to the version without PH-Reg, the improvements are not as striking.

We show in Figure 6, for CLIP family models the artifacts primarily manifest as lower magnitude tokens when utilized for open-vocabulary segmentation. While in DINOv2 the artifacts are high magnitude -- confirming the observation in "Vision Transformers Need Registers" [1].

We are unsure what causes the difference in magnitudes of artifact tokens, and our offline discussions with DINOv2 authors suggests that they are also unsure why artifacts appear in DINOv2 but not DINOv1. One co-author of DINOv2 suggested that the increase in single-color background images (collected from online retail), and the decrease in natural photographs in the training dataset was a likely reason in the artifact tokens.

Regardless of the cause, if artifact tokens are low magnitude, then their influence decreases after averaging. We believe this is the reason of the increased performance of PH-Reg in CLIP family models. For future work, it may make sense to decouple patch direction and magnitude prior to distillation.

We will include a discussion of this in the revision.

[1] Vision Transformers Need Registers (ICLR 2024)

Conclusion

We thank you for providing detailed and thoughtful feedback. Following your suggestions, we have run an evaluation to measure the impact of register counts in different models, the training efficiency of our method versus DVT, and the inference computational cost. We hope our answers have helped clarify your questions, and look forward to any additional discussion.

We are very grateful for the positive assessment you've given to our work!

We would like to again express our appreciation for your suggestions.

As a gentle reminder, for NeurIPS this year -- you may have to edit your review with a "Final Justification". Thank you for taking the time!

We appreciate your excellent suggestions! Below are our responses to specific questions. We look forward to further discussion, and welcome additional comments.

Q1) Why not vertical flips as augmentation?

We avoid vertical flips in our augmentation pipeline because this violates real-world physical constraints. Gravity imposes a consistent vertical orientation in natural scenes (e.g., sky above ground, human faces upright), making vertically flipped images unrealistic. In contrast, horizontal flips preserve vertical orientation while simulating naturally occurring viewpoint variations (e.g., an object facing left/right). Augmentations like horizontal flips and shifts expand the training distribution with plausible variations.

The mmseg package [1] commonly used for dense vision evaluation, similarly defaults to horizontal for RandomFlip. Generally other works in computer vision also avoid using vertical flips as augmentation. For example, the ResNet paper [2] and DeiT III [3] training recipe (utilized by DINOv2 for ImageNet) adopt horizontal flips without vertical flipping in their augmentation pipeline.

We will update our revision to include a paragraph clarifying why we do not utilize vertical flips in our augmentation function.

[1] https://mmcv.readthedocs.io/en/2.x/api/generated/mmcv.transforms.RandomFlip.html

[2] Deep Residual Learning for Image Recognition (CVPR 2016)

[3] DeiT III: Revenge of the ViT (ECCV 2022)

Q2) Ablation over shift ratios

This is a great question. The optimal shifting ratio of 15% was determined using extensive experiments. Here we provide a table where additional shift ratios are explored using NACLIP teacher features (note that only the student network has registers).

We find that a shifting ratio of 15% provides the optimal performance (mIoU, higher is better):

We will add this table to the paper in an upcoming revision.

Q3) Number of registers

We find that the rate of improvement decreases as additional registers are added. This trend is observed over all models. We provide additional results here (median cosine similarity, higher is better), and will include these results in an upcoming revision:

Q4) Proof of convergence

We would like to clarify, that a key advantage of our method over DVT is that we avoid expensive gradient-based neural field optimization. On page 12 we include a proof showing that we obtain the same optimal result in expectation as the neural field stage of DVT.

Let $f_1, ..., f_n \in R^d$ be feature vectors of a single patch from $n$ transformations of an input image. DVT optimizes a neural field using MSE loss gradients, while we take the unweighted mean. We summarize the proof here:

Let $f^{*} = \text{argmin}_f \sum_i^n ||f_i-f||^2_2$

Expanding the objective gives us: $\text{argmin}_f \sum_i^n (f_i^{T}f_i - 2f_i^{T}f + f^{T}f)$

Dropping constant terms that do not change the optimum: $n f^{T}f - 2(\sum_i^nf_{i}^{T})f$

Dividing and multiplying the right side by n: $n f^{T}f - 2n(\sum_i^n\frac{1}{n}f_{i}^{T})f$

Dividing the equation by constant n shows that we obtain the same optimum as DVT: $||f - \frac{1}{n}\sum_i^nf_i||^2_2$ so $f^* = \frac{1}{n}\sum_i^nf_i$

Due to this, our unweighted mean yields the same optimum as DVT. We will highlight this proof in the upcoming revision.

Q5) Evaluation of non-salient objects

We agree that the main paper could benefit from additional visualizations. To clarify, Figure 4 on page 5 of the main paper provides examples of open-vocabulary queries using the CLIP text branch for different objects within the same image.

We indeed evaluate our model on COCO benchmarks in Table 1 on Page 5 of the main paper (COCO-Objects is listed as Object, COCO-Stuff is listed as Stuff). We also evaluate on other widely used segmentation benchmarks like PASCAL VOC, PASCAL Context, CityScape, and ADE20k. We reproduce the results for COCO here (mIoU, higher is better), and will update the naming in an upcoming revision:

Additional segmentation visualization results are provided in Figure 1 to Figure 3 of the supplemental PDF, with heatmap visualizations of multiple objects within the same image provided from Figure 4 to Figure 7 of the supplemental.

Our results demonstrate PH-Reg yields significant qualitative improvements in non-salient object localization.

We appreciate this suggestion. We will update the revision to include additional visualizations in the main paper.

Q6) DINOv2 performance compared to DVT

We would clarify that our PH-Reg method applied to DINOv2 obtains consistent performance improvements compared the unmodified DINOv2 as shown in Table 2 on page 7 of the main paper. We also outperform the DVT denoiser on depth prediction tasks. Our method also has the advantage of not utilizing gradient-based neural field learning as done in DVT, offering significant training time and space usage improvements compared to DVT.

Here we provide training time evaluations on a single GPU for DVT (single GPU is specified in DVT's official code for all steps of feature extraction and neural field learning) and for fairness we also use a single GPU for our method's training. We follow all their official settings and use identical resolution and measure training time for the same ViT model.

In practice, we parallelize training across 4 GPUs, and complete PH-Reg in roughly 1.5 days.

A further advantage is the space utilization during training. DVT requires saving 1.4 terabytes of intermediate neural fields in the form of serialized Instant-NGP files. Our method computes all distillation targets on the fly, and requires no additional space.

In summary, our method is significantly faster and space efficient compared to DVT -- while holding an advantage in depth prediction. We agree that future work should investigate how to further boost the performance in semantic segmentation when using PH-Reg.

Conclusion

We thank the reviewer for the suggestions, and have added additional quantitative ablations over the shifting ratio and number of registers. Following your suggestion, we have also highlighted the optimality proof we previously included in the supplemental. We will also include additional visualizations in the revision showing the localization quality of different objects within the same image, and clarify the datasets that we evaluate on.

We're grateful for your clear and helpful comments. In light of our response and the positive evaluation from other reviewers, we hope you might view our work in a more favorable light. Please feel free to let us know if you have additional questions!