# Learning to Substitute Components for Compositional Generalization

arXiv: https://arxiv.org/abs/2502.20834 · HTML: https://mlsanity.com/paper/2502.20834/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 69 of its 89 references that are in the map.
- First posted on arXiv 28 Feb 2025.

## Builds on (most cited first, 30 of 69)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [Least-to-Most Prompting Enables Complex Reasoning in Large Language Models](https://mlsanity.com/paper/2205.10625/) (2022)
- [Get To The Point: Summarization with Pointer-Generator Networks](https://mlsanity.com/paper/1704.04368/) (2017)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Learning to Compress Prompts with Gist Tokens](https://mlsanity.com/paper/2304.08467/) (2023)
- [Long-context LLMs Struggle with Long In-context Learning](https://mlsanity.com/paper/2404.02060/) (2024)
- [Measuring Compositional Generalization: A Comprehensive Method on Realistic Data](https://mlsanity.com/paper/1912.09713/) (2019)
- [COGS: A Compositional Generalization Challenge Based on Semantic Interpretation](https://mlsanity.com/paper/2010.05465/) (2020)
- [OpenNMT: Open-source Toolkit for Neural Machine Translation](https://mlsanity.com/paper/1709.03815/) (2017)
- [Good-Enough Compositional Data Augmentation](https://mlsanity.com/paper/1904.09545/) (2019)
- [Compositional Generalization and Natural Language Variation: Can a Semantic Parsing Approach Handle Both?](https://mlsanity.com/paper/2010.12725/) (2020)
- [Diverse Demonstrations Improve In-context Compositional Generalization](https://mlsanity.com/paper/2212.06800/) (2022)
- [Improving Text-to-SQL Evaluation Methodology](https://mlsanity.com/paper/1806.09029/) (2018)
- [Data Recombination for Neural Semantic Parsing](https://mlsanity.com/paper/1606.03622/) (2016)

## Built on by (most cited first, 1 of 1)

- [Scaling Reasoning Hop Exposes Weaknesses: Demystifying and Improving Hop Generalization in Large Language Models](https://mlsanity.com/paper/2601.21214/) (2026)

## Abstract (arXiv)

Despite the rising prevalence of neural language models, recent empirical evidence suggests their deficiency in compositional generalization. One of the current de-facto solutions to this problem is compositional data augmentation, which aims to introduce additional compositional inductive bias. However, existing handcrafted augmentation strategies offer limited improvement when systematic generalization of neural language models requires multi-grained compositional bias (i.e., not limited to either lexical or structural biases alone) or when training sentences have an imbalanced difficulty distribution. To address these challenges, we first propose a novel compositional augmentation strategy called Component Substitution (CompSub), which enables multi-grained composition of substantial substructures across the entire training set. Furthermore, we introduce the Learning Component Substitution (LCS) framework. This framework empowers the learning of component substitution probabilities in CompSub in an end-to-end manner by maximizing the loss of neural language models, thereby prioritizing challenging compositions with elusive concepts and novel contexts. We extend the key ideas of CompSub and LCS to the recently emerging in-context learning scenarios of pre-trained large language models (LLMs), proposing the LCS-ICL algorithm to enhance the few-shot compositional generalization of state-of-the-art (SOTA) LLMs. Theoretically, we provide insights into why applying our algorithms to language models can improve compositional generalization performance. Empirically, our results on four standard compositional generalization benchmarks(SCAN, COGS, GeoQuery, and COGS-QL) demonstrate the superiority of CompSub, LCS, and LCS-ICL, with improvements of up to 66.5%, 10.3%, 1.4%, and 8.8%, respectively.