From independent patches to coordinated attention: Controlling information flow in vision transformers

TL;DR

This paper introduces a method to explicitly measure and control information flow in vision transformers by applying variational information bottlenecks, enabling a spectrum from independent patch processing to full global attention, and analyzing how representations evolve.

Contribution

The authors propose a novel approach to control and measure information flow in vision transformers using variational bottlenecks, facilitating interpretability and control over attention mechanisms.

Findings

Models with constrained information flow are more interpretable.

Global visual representations emerge from local patch processing.

Controlled models show different classification behaviors.

Abstract

We make the information transmitted by attention an explicit, measurable quantity in vision transformers. By inserting variational information bottlenecks on all attention-mediated writes to the residual stream -- without other architectural changes -- we train models with an explicit information cost and obtain a controllable spectrum from independent patch processing to fully expressive global attention. On ImageNet-100, we characterize how classification behavior and information routing evolve across this spectrum, and provide initial insights into how global visual representations emerge from local patch processing by analyzing the first attention heads that transmit information. By biasing learning toward solutions with constrained internal communication, our approach yields models that are more tractable for mechanistic analysis and more amenable to control.