Towards Compositional Generalization of LLMs via Skill Taxonomy Guided Data Synthesis

TL;DR

The paper tackles the data sparsity barrier in compositional generalization by introducing STEPS, a framework that learns a hierarchical Skill Taxonomy through structural entropy minimization on a skill co-occurrence graph and then synthesizes data by maximizing marginal structural information under hierarchical constraints. It combines a bottom-up taxonomy induction with a recursive, entropy-guided data synthesis approach to generate coherent, high-information skill compositions. Empirical results across MT-Bench, AlpacaEval 2.0, and WildBench show STEPS improves performance over baselines and demonstrates favorable scaling across compositional complexity and data volume, with evidence from taxonomy-guided curriculum learning and agentic benchmarks. The work provides a principled, scalable method to enhance compositional generalization in LLMs and agentic systems, enabling more robust handling of complex, multi-skill tasks.

Abstract

Large Language Models (LLMs) and agent-based systems often struggle with compositional generalization due to a data bottleneck in which complex skill combinations follow a long-tailed, power-law distribution, limiting both instruction-following performance and generalization in agent-centric tasks. To address this challenge, we propose STEPS, a Skill Taxonomy guided Entropy-based Post-training data Synthesis framework for generating compositionally challenging data. STEPS explicitly targets compositional generalization by uncovering latent relationships among skills and organizing them into an interpretable, hierarchical skill taxonomy using structural information theory. Building on this taxonomy, we formulate data synthesis as a constrained information maximization problem, selecting skill combinations that maximize marginal structural information within the hierarchy while preserving semantic coherence. Experiments on challenging instruction-following benchmarks show that STEPS outperforms existing data synthesis baselines, while also yielding improved compositional generalization in downstream agent-based evaluations.

Figures (7)

• Figure 1: Illustration of the STEPS framework.
• Figure 2: Impact of compositional complexity $k$.
• Figure 3: Impact of data size on Llama-3-8B.
• Figure 4: Comparative performance on WB Score across different paradigms. We evaluate Qwen2.5 and Llama3 models in both (a) Base and (b) Instruct settings.
• Figure 5: The system prompt used by STEPS.