Vision Transformers Need Cross Aggregation for Patch Semantics

Qualitative results

Patch features on random Cityscapes frames

Six val frames drawn at random with a fixed seed, not selected. Pick a pre-training and a frame.

Top three principal components of the frozen last-layer patch features, as RGB, on a 1024×1024 crop (64×64 patches). PCA is fitted per model and per image, so colours carry no meaning across panels: compare the structure, such as poles, pedestrians and lane markings.

Method

CALiT changes only the FFN. After fc1 and GELU, each token's 4D-wide hidden is read as S slots of concepts, and two light modules act on it in parallel before fc2.

• Cross Aggregation (CAg). The CLS slots attend over the slots of all patches: CLS slots are the queries, patch slots both the keys and the values, with as many heads as self-attention. There are no projection weights, so the operation is parameter-free and adds 0.6 % FLOPs to ViT-S. A loss on the CLS token now flows back to every patch through this softmax, besides the self-attention route.
• Local Mixing (LMix). A 3×3 depthwise convolution over the patch hiddens on their √N×√N grid, one kernel per channel and no extra nonlinearity. It is Dirac-initialised, so at step 0 the block computes cross aggregation alone. A depthwise convolution alone, as in LocalViT, lowers the semantic probes on 3 of 4 datasets; added to CAg, it raises all four.
• Cost. +0.75 % parameters and +1.44 % FLOPs for ViT-S, +0.40 % and +1.03 % for ViT-B. S = 4 slots for ViT-S and 8 for ViT-B (head size 64 and 32).

# x: (1+N, D) tokens, CLS first, patches on a √N×√N grid.  S slots, H heads (as in SA).
y       = x + SA(LN(x))
h       = GELU(W1(LN(y)))              # (1+N, 4D)
h_c     = h[:1];  h_p = h[1:]          # split the CLS row from the N patch rows

# cross aggregation: parameter-free, MHA without projection weights
C       = h_c.reshape(S, 4D/S)         # CLS slots: queries
P       = h_p.reshape(N*S, 4D/S)       # patch slots: keys and values (un-mixed h_p)
h_c_new = MHA_H(C, P, P).reshape(1, 4D)

# local mixing: depthwise 3×3, Dirac-initialised, patches only
G       = h_p.T.reshape(4D, √N, √N)
h_p_new = DWConv3x3(G).reshape(4D, N).T

z       = y + W2(concat[h_c_new, h_p_new])
return z

Patch gradients, measured

We take the gradient of each model's own CLS loss with respect to the 196 patch tokens entering the last block, on all 38,651 ImageNet-val images whose object box leaves at least one patch outside. Two questions: how many directions does that gradient span, and does it single out the object?

Structure of the patch gradient at the input of the last block. Effective rank: entropy-effective rank of the 196×D gradient matrix; a plain ViT is bounded by 2n_h (two routes per head). Contrast: agreement with the mean in-box gradient direction, inside minus outside the box, each patch weighted by its gradient norm. Last column: share of images on which CALiT's contrast exceeds ViT's; its effective rank exceeds ViT's on ≥ 99.9 % of images.

Gradient agreement, per model

• ViT's patch gradients agree with the object direction inside and outside the box, and their effective rank stays at or below the 2n_h bound (12 for S, 24 for B).
• CALiT's effective rank is 2.6–3.6× ViT's, and its norm-weighted contrast is 1.6–2.6× higher, on 77–91 % of images.
• In the supervised models the rank gap is there from epoch 100 and stays constant, while the contrast gap widens with training. Cross aggregation alone (CAT) already lifts the rank; local mixing adds contrast.

Results

Semantic segmentation with a frozen backbone (Table 1 of the paper). DINOv2-style linear probe, a BatchNorm and a 1×1 convolution on the last-layer patch tokens, mIoU ↑. IN-1k top-1 is measured as each baseline does: supervised accuracy, a linear probe (lp) for self-supervised models, zero-shot (zs) after vision-language fine-tuning. Δ: CALiT minus the ViT of the same block. * mIoU as reported by Marouani et al. with the same probe. Bold: best per block. CALiT gives up a little ImageNet top-1 in 5 of the 7 settings, by 2.8 points at most (DINOv2-B).

Thin structures

The largest gains are on thin classes: fence, pole, traffic light, traffic sign, rider, motorcycle and bicycle. We also probe ACDC, urban driving under fog, rain and snow, with the Cityscapes head (zero-shot) and with a head trained on ACDC.

Thin-structure segmentation and robustness (mIoU ↑, linear probe as above). thin: the 7 thin classes. ACDC zero-shot: the Cityscapes-trained head applied to ACDC val. ACDC trained: a head trained from scratch on the 1600 ACDC training images for 4k iterations, evaluated on the 400 val images.

Ablations

Single-variable ablations of CALiT-S (supervised IN-1k, linear probes as above; bold best, italic second). Placed only in the last block, CAg loses most of the dense gain: the earlier blocks still get the narrow gradient, so CAg belongs in every block. Applying CAg to half of the FFN width recovers most of the gain. A uniform mean in place of attention loses patch quality, more than half-width CAg does.

BibTeX

@article{chhatkuli2026calit,
  title   = {Vision Transformers Need Cross Aggregation for Patch Semantics},
  author  = {Chhatkuli, Ajad and Paudel, Pramish and Zhang, Deheng and
             Wu, Kanzhi and Van Gool, Luc and Paudel, Danda Pani},
  journal = {arXiv preprint},
  year    = {2026}
}