Training Vision Transformers to Focus: Emergence of Spatial Attention Head Specialization

Training Vision Transformers to Focus. Emergence of Spatial Attention Head Specialization. By Yuji Nozawa and two others, from Kioxia. Published on the ninth of November.

Self-supervised Vision Transformer models such as DINO learn robust visual perception from large image datasets, with no labels required. What they do not have is a focusing mechanism during training. Without one, the perceptual representations they form are often ambiguous, or too similar to one another to tell apart.

The proposal is to give training that missing focus. Two visual focusing loss functions are introduced, each designed to tie the model's attention to specific regions within an image, so that where the model looks corresponds to something definite in the picture.

The lever is the model's own multi-head attention. The losses use it to steer the model's focus selectively, and what emerges is spatial attention head specialization: heads that become diverse and spatially focused in what they perceive. None of this calls for annotations. The specialization arises from the training objective rather than from labels.

Evaluation is both qualitative and quantitative. The method substantially increases how spatially selective the attention heads are, and how diverse they are from one another. Those gains carry through to interpretability, and to performance on dense prediction tasks, with zero-shot image segmentation given as the example.

Heads that once tended toward overlapping views of an image come out of training each attending to its own region. That division of labour is the specialization of the title, and it is reached through the loss alone, without anyone marking which region belongs to which head.