Dinov2 With Registers includes register tokens in patch token mean for classifier input
Author: psandovalsegura
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.48.2
- Platform: Linux-4.18.0-553.50.1.el8_10.x86_64-x86_64-with-glibc2.28
- Python version: 3.9.21
- Huggingface_hub version: 0.28.1
- Safetensors version: 0.5.2
- Accelerate version: 1.3.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.6.0+cu124 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: no

### Decription

In the `Dinov2WithRegistersForImageClassification` model, the classifier input is constructed by concatenating the [CLS] token with the mean of all image _patch_ tokens. But currently, this mean includes register tokens, which are intended to capture global information and are not representative of specific image patches.

**Current Implementation Reference**:
https://github.com/huggingface/transformers/blob/397a5ede33863d6f7137c771a68d40036cac0396/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py#L773-L778

**Expected Implementation**: The mean should be computed only over the patch tokens. We need to skip the register tokens which are located after the [CLS] token, as follows:
```
sequence_output = outputs[0]  # batch_size, sequence_length, hidden_size

cls_token = sequence_output[:, 0]
patch_tokens = sequence_output[:, 1+self.config.num_register_tokens:] # change here!

linear_input = torch.cat([cls_token, patch_tokens.mean(dim=1)], dim=1)
```

This error is because we forgot to discard the registers at the end of the forward pass, as is done in the original paper implementation here:
https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/hub/classifiers.py#L61-L68

@NielsRogge @BernardZach @qubvel @ArthurZucker 


### Reproduction

**Without registers**
Using `AutoModelForImageClassification.from_pretrained('facebook/dinov2-small-imagenet1k-1-layer')`: For an image with 256 image patches and with a breakpoint at line 773 of [transformers/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py](https://github.com/huggingface/transformers/blob/397a5ede33863d6f7137c771a68d40036cac0396/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py#L773-L778), we get the following shapes:
```
sequence_output.shape=torch.Size([128, 257, 384])
cls_token.shape=torch.Size([128, 384])
patch_tokens.shape=torch.Size([128, 256, 384])
```
which are all as expected.

**With registers**
Using `AutoModelForImageClassification.from_pretrained('facebook/dinov2-with-registers-small-imagenet1k-1-layer')`: For an image with 256 image patches and with a breakpoint at line 773 of [transformers/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py](https://github.com/huggingface/transformers/blob/397a5ede33863d6f7137c771a68d40036cac0396/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py#L773-L778), we get the following shapes:
```
sequence_output.shape=torch.Size([128, 261, 384]) # includes the 4 extra registers in sequence dim, as expected
cls_token.shape=torch.Size([128, 384])            # as expected
patch_tokens.shape=torch.Size([128, 260, 384])    # should be 256 along sequence dim, not 260
```
but here we see that `patch_tokens` still includes the registers.

### Expected behavior

**With registers**
Using `AutoModelForImageClassification.from_pretrained('facebook/dinov2-with-registers-small-imagenet1k-1-layer')`: For an image with 256 image patches and with a breakpoint at line 773 of [transformers/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py](https://github.com/huggingface/transformers/blob/397a5ede33863d6f7137c771a68d40036cac0396/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py#L773-L778), **should have** the following shapes:
```
sequence_output.shape=torch.Size([128, 261, 384])
cls_token.shape=torch.Size([128, 384])                    
patch_tokens.shape=torch.Size([128, 256, 384])  
```

Comments:
Thanks for the ping, that's a mistake on my side.

Given they [all have < 1k downloads](https://huggingface.co/models?sort=trending&search=facebook%2Fdinov2-with-registers), I assume we can reupload them with the appropriate fix. 

cc @qubvel 



Thanks for reporting the issue, @psandovalsegura!
I agree with @NielsRogge, it would be best to update the model.

@psandovalsegura, would you be interested in implementing the fix? If so, here’s the plan:

- Update the modeling code with the fix.
- Reconvert the classification models and open PRs to the corresponding HF repos for ImageNet-1K models ([list of repos](https://huggingface.co/models?sort=trending&search=facebook%2Fdinov2-with-registers)).

Thanks so much for your help!

@NielsRogge Why is there a need to reconvert the checkpoints as the fix is in forward pass and its simply exclude register tokens form path tokens calculation 🤔 

@qubvel I could try to fix it! My main question is how do I regenerate [models/dinov2_with_registers/modeling_dinov2_with_registers.py](https://github.com/huggingface/transformers/blob/397a5ede33863d6f7137c771a68d40036cac0396/src/transformers/models/dinov2_with_registers/modeling_dinov2_with_registers.py#L773-L778)? The header on this file says not to modify it directly.

The only change I'd make (within `Dinov2WithRegistersForImageClassification.forward`) is:
`patch_tokens = sequence_output[:, 1+self.config.num_register_tokens:]`

I took a look at the [original PR](https://github.com/huggingface/transformers/pull/35348) for Dinov2 with registers but can't really tell what the steps were to do it.

Yes you're right, sorry no need to convert.

In that case we only need to modify [this class](https://github.com/huggingface/transformers/blob/397a5ede33863d6f7137c771a68d40036cac0396/src/transformers/models/dinov2_with_registers/modular_dinov2_with_registers.py#L316). The modular file can use inheritance, hence we can inherit from `Dinov2ForImageClassification` but change the forward pass from [here](https://github.com/huggingface/transformers/blob/9c5b1319d01d8a7455a0de0ff8f7b49064d26624/src/transformers/models/dinov2/modeling_dinov2.py#L732) to discard the register tokens.  

After editing the modular file, you can run this from the root of the repo:

```bash
python utils/modular_model_converter.py --files_to_parse src/transformers/models/dinov2_with_registers/modular_dinov2_with_registers.py
```
which is going to update modeling_dinov2_with_registers.py.

Read more here: https://huggingface.co/docs/transformers/v4.51.3/en/modular_transformers#modular-transformers

 @psandovalsegura I went ahead and raised a fix PR — hope you don't mind!. I am also working on the Dinov2 for my task and thanks for the correction 🤗