🤖 machine learning

Compositional Generalization in Autoregressive Models via Logit Composition

This paper introduces a principled logit composition strategy for autoregressive models, inspired by diffusion methods, which ensures projective control over output subspaces and preserves length generalization under factorized-conditionals assumptions.

Original authors: Aakash Kumar, Maria Sofia Bucarelli, Emanuele Natale

Published 2026-05-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Aakash Kumar, Maria Sofia Bucarelli, Emanuele Natale

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a team of three very specific chefs working in a kitchen.

Chef Base (The Background): This chef is great at following the recipe exactly as written. If you ask for a plain sandwich, they make a perfect plain sandwich. They don't add anything extra.
Chef Math: This chef is a wizard at adding cheese and spices, but only if the recipe calls for it. If the recipe says "add cheese," they do it perfectly. If it doesn't, they leave the sandwich alone.
Chef Code: This chef is an expert at adding lettuce and tomatoes, but only when the recipe asks for it.

The Problem:
Usually, if you want a sandwich with both cheese and lettuce, you have to hire one giant chef who knows how to do everything at once. Or, you try to mix the three chefs together by blending their aprons (weights) or telling them to take turns (routing). But often, this causes chaos: Chef Math might try to add cheese when Chef Code is trying to add lettuce, or they might argue over how to finish the sandwich, resulting in a messy, inedible pile of food.

The Paper's Solution: "Logit Composition"
The authors of this paper propose a new way to combine these chefs. Instead of blending their aprons or making them take turns, they let all three chefs shout out their suggestions for the next ingredient at the same time, but with a special rule:

The Rule: They take Chef Math's suggestion, add Chef Code's suggestion, and then subtract Chef Base's suggestion.
The Magic: Because Chef Base is just the "plain" version, subtracting their voice cancels out the noise.
- When the recipe needs cheese, Chef Math is loud, Chef Code is quiet (because they don't know about cheese), and Chef Base is neutral. The math works out so that the final decision is "Add Cheese."
- When the recipe needs lettuce, Chef Code is loud, Chef Math is quiet, and Chef Base is neutral. The final decision is "Add Lettuce."
- When the recipe needs nothing special, everyone is quiet, and the sandwich stays plain.

Why This Works (The "Disjoint" Secret)
The paper argues that this only works perfectly if the chefs have disjoint jobs.

Chef Math only touches the "cheese steps."
Chef Code only touches the "lettuce steps."
They never try to do the other's job at the same time.

If Chef Math tries to add cheese while Chef Code is adding lettuce, they might crash into each other. But if they know exactly when to act (e.g., "I only act on step 3, you only act on step 5"), they can work together seamlessly without fighting. The paper calls this "Factorized Conditionals."

The "Projective" Guarantee
The paper proves that if the chefs stick to their specific zones, the final sandwich (the output) will be exactly what you'd get if you had hired a super-chef who knew both skills perfectly.

The "Cheese" part of the sandwich comes 100% from Chef Math.
The "Lettuce" part comes 100% from Chef Code.
The "Bread" part comes 100% from Chef Base.

They don't interfere with each other. It's like a projection: if you look at just the cheese part, it looks exactly like Chef Math made it. If you look at the lettuce, it looks exactly like Chef Code made it.

Does it work for long recipes? (Length Generalization)
The paper also checked if this works for very long recipes (like a 100-page cookbook) that the chefs haven't seen before. They found that as long as the chefs know their specific "zones" (e.g., "I handle steps 10–20") and the recipe follows the same pattern, the combined team can handle the long recipe just as well as the short one. They don't get confused just because the book is thicker.

Real-World Test
The authors tested this with real AI models (Large Language Models):

They took a base model (Gemma 2).
They took a version fine-tuned for Math.
They took a version fine-tuned for Coding.
They combined them using their "shout-out" rule.

The Result:
The new combined model got better at coding than the coding expert alone, and it kept almost all of the math expert's skills. It didn't lose its math ability just because it learned to code. It was a "best of both worlds" situation without needing to retrain the whole team from scratch.

The Catch
The paper admits this works best when the tasks are clearly separated. If you try to combine too many experts (like adding a "Poetry" chef and a "History" chef to the mix), the team might get confused about when to stop talking, leading to gibberish. But for clear, separate skills, this "Logit Composition" is a powerful, principled way to build modular AI systems.

Technical Summary: Compositional Generalization in Autoregressive Models via Logit Composition

Problem Statement

Composing specialized generative models to create systems that are modular, reusable, and controllable remains a core challenge in large language model (LLM) research. Existing approaches for autoregressive models typically operate in parameter space (e.g., weight interpolation, task arithmetic, sparsity-based merging) or architecture space (e.g., Mixture-of-Experts with dynamic routing). While these methods combine fine-tuned checkpoints, they often lack explicit, static guarantees regarding the resulting output distribution or fail to provide a principled rule for how components interact without interference.

Conversely, diffusion models have demonstrated success in compositional generation by linearly combining score functions, a perspective recently formalized by Bradley et al. [2025] through the concept of projective composition. However, this framework has not been extended to autoregressive models, where generation is inherently sequential and conditioned on an evolving prefix, making the applicability of specific skills dependent on the current state of the generated sequence.

Methodology

The authors propose a selective composition strategy for autoregressive systems (such as Transformers) that operates directly in probability space rather than parameter space. Given a background model $p_b$ and a set of specialized models $\{p_1, \dots, p_k\}$ , the composed model $\hat{p}$ is defined by an explicit arithmetic rule applied to conditional next-token distributions:

$\hat{p}(y_t | x, y_{<t}) = \frac{1}{Z_t} p_b(y_t | x, y_{<t}) \prod_{i=1}^k \frac{p_i(y_t | x, y_{<t})}{p_b(y_t | x, y_{<t})}$

where $Z_t$ is a normalizing constant. This formulation effectively adds the log-probabilities of the experts and subtracts the log-probabilities of the background model.

Theoretical Framework

The paper establishes conditions under which this composition is provably correct and preserves the distinct behaviors of component models:

Factorized Conditionals (Token Space): The authors define a property where the generation space is partitioned into disjoint sets of token indices (masks). If each specialized model $p_i$ acts independently on its designated mask $M_i$ and coincides with the background model $p_b$ elsewhere, the composition is projective. This means the marginal distribution of the composed model on any mask $M_i$ is identical to that of the specialized model $p_i$ .
Skill Factorizability (Sequential Setting): Recognizing that autoregressive tasks (like trigonometric proofs) depend on the history of generated tokens, the authors introduce skill factorizability. A model is skill factorizable if its behavior can be decomposed into a background phase, a specific skill phase (defined by stopping times $\tau_1, \tau_2$ ), and a return to background.
Joint Factorizability: For a collection of models to compose correctly, their active intervals (where they deviate from the background) must be pairwise disjoint. Under this condition, the composed model applies each skill exactly during its designated interval and reverts to the background model otherwise.
Feature Space Extension: The framework is extended to abstract feature spaces via an invertible map $D$ . If the pushforward distributions satisfy the factorized conditionals property in feature space, the composition remains valid, allowing for the combination of models that may not have disjoint token domains but disjoint "conceptual" domains.
Length Generalization: The paper argues that if the factorization assumptions and component guarantees hold uniformly at a target length, the composed model inherits the length-generalizing behavior of its components.

Connection to Model Merging

The authors draw a theoretical link between their probability-space composition and Task Arithmetic (Ilharco et al. [2023]). They demonstrate that in a linear setting, adding task vectors in parameter space is mathematically equivalent to the multiplicative composition of output distributions. This suggests that the empirical success of model merging techniques may be explained by an underlying additive composition principle where task updates affect disjoint regions of the output space.

Key Contributions

Principled Framework: Introduction of a probability-space framework for composing autoregressive models, extending the concept of projective composition from diffusion models to sequential generation.
Correctness Guarantees: Identification of specific conditions (factorized conditionals and joint factorizability) under which composition is provably correct, ensuring that components preserve control over their designated subspaces without interference.
Length Generalization: Demonstration that composition preserves length-generalizing behavior when factorization assumptions hold uniformly at the target length.
Theoretical Bridge: Connection of the proposed framework to existing model merging methods (specifically Task Arithmetic), providing a theoretical explanation for why certain merging heuristics succeed.
Empirical Validation: Verification of the approach on both synthetic tasks and large language model benchmarks.

Results

The authors validate their approach through two primary experimental settings:

Controlled Toy Task: In a letter-replacement task where models were trained to replace specific characters (e.g., 'A' $\to$ 'B', 'C' $\to$ 'D'), the composed model achieved 100% success in performing both transformations simultaneously. This confirmed that the composition operator correctly integrates disjoint skills.
LLM Benchmarks: Using a Gemma 2 2B base model ( $p_{base}$ $p_{ba se}$ ) and two fine-tuned experts ( $p_{math}$ $p_{ma t h}$ and $p_{coding}$ $p_{co d in g}$ ), the authors composed a model $\hat{p}$ $\overset{p}{^}$ and evaluated it on four benchmarks:
- Coding: On HumanEval+ and MBPP+, the composed model outperformed the coding expert alone (30.5 vs. 24.4 and 39.4 vs. 38.1, respectively).
- Mathematics: On GSM8K and MATH, the composed model retained significant performance from the math expert (41.2 vs. 50.2 and 24.2 vs. 25.5), demonstrating that the coding skill did not catastrophically interfere with mathematical reasoning.
- Limitation: The authors note that composing more than two experts led to a sharp degradation in generation quality, likely due to instability in end-of-sequence token prediction and termination behavior.

Significance and Claims

The paper claims to provide a principled understanding of when and why model composition and merging succeed in autoregressive systems. By shifting the focus from parameter space to probability space, the authors identify conditions under which interactions between models remain stable and non-interfering.

The significance lies in offering a static, explicit composition rule with theoretical guarantees, contrasting with dynamic routing methods or heuristic parameter averaging. The work suggests that autoregressive model composition can be understood as a structured probabilistic operation, potentially enabling the construction of more modular and controllable LLMs. However, the authors remain modest, acknowledging that their factorization assumptions may not fully hold in practice for all pre-trained models and that the empirical evaluation is currently limited in scope.