Transformers as Communication Systems: Controlling Information Flow with Bottlenecks
Abstract
We make the information communicated by attention between residual streams in vision transformers a measurable and controllable quantity. 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 spectrum that interpolates between independent patch processing and fully expressive global attention. On ImageNet-100, we characterize how classification and self-supervised representation learning change across this spectrum, revealing how information flow is allocated across depth, heads, and patches as global visual representations emerge from local processing. We further inspect the first attention heads that transmit information, identifying simple visual computations that appear under tight communication constraints. 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.