A Sharper Picture of Generalization in Transformers

TL;DR

This paper investigates the generalization behavior of transformers on boolean domains through Fourier spectra analysis, proposing PAC-Bayes bounds based on spectral sparsity and validating findings empirically.

Contribution

It introduces a novel PAC-Bayes approach to bound transformer generalization using spectral sparsity, contrasting prior Rademacher complexity methods.

Findings

Sparse low-degree spectra enable good generalization bounds.

Existence of flat minima for boolean functions supports theoretical bounds.

Empirical validation confirms the theoretical predictions.

Abstract

We study transformers' generalization behavior on boolean domains from the perspective of the Fourier Spectra of their target functions. In contrast to prior work (Edelman et al., 2022; Trauger and Tewari, 2024), which derived generalization bounds from Rademacher complexity, we investigate the feasibility of obtaining generalization bounds via PAC-Bayes theory. We show that sparse spectra concentrated on low-degree components enable low-sharpness constructions with good generalization properties. Our idea is to show the existence of flat minima implementing any boolean function of sparsity no greater than the context length, and then apply a PAC-Bayes bound to an idealized low-sharpness learner, resulting in a non-vacuous generalization bound. We evaluate predictions empirically and conduct a mechanistic interpretability study to support the realism of our theoretical construction in real transformers.