Explaining Grokking in Transformers through the Lens of Inductive Bias

TL;DR

This paper explores the phenomenon of grokking in transformers by examining how architectural choices and optimization settings influence learning dynamics, revealing the role of inductive bias and feature compressibility in generalization.

Contribution

It provides a detailed analysis of how inductive biases from architecture and optimization affect grokking, highlighting the nuanced and continuous nature of feature evolution in transformers.

Findings

Layer Normalization position significantly affects grokking speed.

Optimization parameters like learning rate and weight decay influence grokking interpretations.

Feature compressibility correlates with generalization in grokking.

Abstract

We investigate grokking in transformers through the lens of inductive bias: dispositions arising from architecture or optimization that let the network prefer one solution over another. We first show that architectural choices such as the position of Layer Normalization (LN) strongly modulates grokking speed. This modulation is explained by isolating how LN on specific pathways shapes shortcut-learning and attention entropy. Subsequently, we study how different optimization settings modulate grokking, inducing distinct interpretations of previously proposed controls such as readout scale. Particularly, we find that using readout scale as a control for lazy training can be confounded by learning rate and weight decay in our setting. Accordingly, we show that features evolve continuously throughout training, suggesting grokking in transformers can be more nuanced than a lazy-to-rich transition of the learning regime. Finally, we show how generalization predictably emerges with feature compressibility in grokking, across different modulators of inductive bias. Our code is released at https://tinyurl.com/y52u3cad.