Bias-driven Alignment of Linear and ReLU Networks

Jirko Rubruck, Kai Jappe Sandbrink, Christopher Summerfield, Andrew M Saxe, Devon Jarvis

Track: Extended Abstract Track

Keywords: Learning dynamics, Representational alignment

TL;DR: When we equip ReLU networks with bias terms they learn similar to linear networks.

Abstract: ReLU networks and their variants are a key building block of modern deep learning architectures. Despite their ubiquity, our understanding of learning dynamics in these models is still limited. Previous work has relied on a strong set of simplifying assumptions such as the removal of bias terms or predefined gating structures. Here, we explore empirically how the inclusion of bias terms influences learning dynamics in ReLU networks in the rich learning regime. Surprisingly, we find that the inclusion of bias terms simplifies learning dynamics, i.e. ReLU networks with bias terms have learning dynamics that are strongly aligned to those of well-understood linear models. Further, ReLU and linear networks with bias terms trained on nonlinear problems display a transient correspondence early in learning that is also reflected in highly structured, linear-like representations. We also highlight additional downstream effects of early linearity and find that the inclusion of bias terms boosts simplicity biases and the over-representation of features associated with simple tasks. We demonstrate the practical relevance of our results beyond simplified settings and show that bias terms can also induce early linearity on image classification tasks. Our results illustrate how seemingly minor and common architectural choices can change learning dynamics, biases towards simplicity, and representational alignment between systems.

Submission Number: 83