@sinahmr this is fairly redundant when the attention pool is already there as an option isn't it?

Thanks for the feedback @rwightman.

Conceptually, PRR is not intended to improve the pooled [CLS] representation itself. Its purpose is to refine the final-layer patch tokens by allowing the classification signal to propagate to spatial tokens in a more diverse, content-dependent way during pretraining.

We discuss this in the paper in comparison to GAP as well. While GAP does route supervision to patch tokens, unlike when global_pool='token' which mainly supervises the [CLS] path, it does so uniformly across spatial locations. For dense prediction, that uniform gradient flow is not ideal. PRR instead redistributes information across tokens non-uniformly before the head. More details are provided in Section 4.2, and the empirical comparison to GAP is shown in Table 5, where PRR yields clear segmentation gains.

On the implementation side, I agree that the current integration may not be ideal. If it would be better represented as part of the pooling path, I can rework it that way, for example by treating it as an option for global_pool. Please let me know if that would fit better.

If the concern is that this is too paper-specific to justify a dedicated switch in vision_transformer.py, I understand that as well. In that case, I would be happy to close the PR.

@sinahmr could you take a look at #2685 ... does it achieve the same end goal? I don't want to introduce high risk changes here so I feel using as an alternate pooling mechanism is more appropriate

Changes in #2685 look good to me, thank you for the implementation!
I'll close this PR.