Making Transformers Logical
Transformers often generalize poorly by defaulting to simple patterns instead of complex functions. The authors tested whether injecting explicit structural priors into attention could fix this. They found that initializing an additive attention mask directly, based on task structure rather than query-key similarity, allows the prior to persist through training. This method significantly improves extrapolation on Boolean reasoning tasks and maintains strong performance across arithmetic and vision benchmarks.