On Provable Length and Compositional Generalization
This addresses the problem of out-of-distribution generalization for researchers and practitioners in machine learning, offering theoretical insights into model robustness, though it is incremental in extending existing theoretical frameworks.
The paper provides the first provable guarantees for length and compositional generalization in common sequence-to-sequence models, showing that limited capacity variants achieve these capabilities with diverse training data, and that chain-of-thought supervision enables length generalization in higher capacity models.
Out-of-distribution generalization capabilities of sequence-to-sequence models can be studied from the lens of two crucial forms of generalization: length generalization -- the ability to generalize to longer sequences than ones seen during training, and compositional generalization: the ability to generalize to token combinations not seen during training. In this work, we provide first provable guarantees on length and compositional generalization for common sequence-to-sequence models -- deep sets, transformers, state space models, and recurrent neural nets -- trained to minimize the prediction error. We show that \emph{limited capacity} versions of these different architectures achieve both length and compositional generalization provided the training distribution is sufficiently diverse. In the first part, we study structured limited capacity variants of different architectures and arrive at the generalization guarantees with limited diversity requirements on the training distribution. In the second part, we study limited capacity variants with less structural assumptions and arrive at generalization guarantees but with more diversity requirements on the training distribution. Further, we also show that chain-of-thought supervision enables length generalization in higher capacity counterparts of the different architectures we study.