How Do Language Models Compose Functions? (2510.01685v1)

Abstract: While LLMs appear to be increasingly capable of solving compositional tasks, it is an open question whether they do so using compositional mechanisms. In this work, we investigate how feedforward LLMs solve two-hop factual recall tasks, which can be expressed compositionally as $g(f(x))$. We first confirm that modern LLMs continue to suffer from the "compositionality gap": i.e. their ability to compute both $z = f(x)$ and $y = g(z)$ does not entail their ability to compute the composition $y = g(f(x))$. Then, using logit lens on their residual stream activations, we identify two processing mechanisms, one which solves tasks $\textit{compositionally}$, computing $f(x)$ along the way to computing $g(f(x))$, and one which solves them $\textit{directly}$, without any detectable signature of the intermediate variable $f(x)$. Finally, we find that which mechanism is employed appears to be related to the embedding space geometry, with the idiomatic mechanism being dominant in cases where there exists a linear mapping from $x$ to $g(f(x))$ in the embedding spaces. We fully release our data and code at: https://github.com/apoorvkh/composing-functions .

• The paper demonstrates that LLMs exhibit a compositionality gap, solving individual sub-tasks but often failing to compose functions reliably.
• It distinguishes processing mechanisms by contrasting explicit intermediate representations with direct input-output mapping without detectable sub-components.
• The study reveals that embedding space linearity strongly influences mechanism selection, impacting compositional generalization even as model size increases.

Mechanisms of Function Composition in LLMs

Introduction

This paper investigates the internal mechanisms by which LLMs perform compositional reasoning, specifically focusing on two-hop factual recall tasks expressible as $y = g(f(x))$. The central question is whether LLMs solve such tasks by explicitly computing intermediate variables (i.e., compositional processing) or by directly mapping inputs to outputs without intermediate representations (i.e., idiomatic processing). The paper extends prior work on the "compositionality gap" and leverages interpretability techniques such as logit lens and token identity patchscope to analyze model activations.

Compositionality Gap in Modern LLMs

The compositionality gap refers to the phenomenon where LLMs can solve individual sub-tasks ($x \to f(x)$ and $f(x) \to g(f(x))$) but fail to solve the overall composition ($x \to g(f(x))$) reliably. The paper confirms that this gap persists in contemporary models, including Llama 3 and OLMo 2, across a diverse set of tasks (arithmetic, factual recall, lexical functions, translation, etc.). The gap is quantified as the proportion of cases where all hops are solved correctly but the composition fails.

Figure 1: Compositionality gap for Llama 3 (3B) across tasks, showing the proportion of examples where all hops are solved but the composition is not.

Scaling model size reduces the gap monotonically but with diminishing returns; for example, Llama 3 (1B $\to$ 405B) sees the gap drop from 72% to 39%, but the gap remains substantial even in the largest models. Instruction tuning and reasoning models (e.g., DeepSeek-R1, o4-mini) further reduce the gap, but do not eliminate it.

Figure 2: Compositionality gap across model families and sizes, with error bands showing interquartile range.

Analysis of Processing Mechanisms

The paper employs logit lens and token identity patchscope to decode intermediate representations from the residual stream. Two distinct processing mechanisms are identified:

• Compositional Mechanism: The model surfaces representations for $x$, then $f(x)$, and finally $g(f(x))$, indicating explicit computation of the intermediate variable.
• Direct Mechanism: The model maps $x$ directly to $g(f(x))$ without a detectable representation of $f(x)$.

The presence of intermediate variables is measured via reciprocal rank in the vocabulary space. Notably, both mechanisms can yield correct answers, and the choice of mechanism is only weakly correlated with accuracy.

Figure 3: Aggregate processing signatures for tasks where Llama 3 (3B) solves all hops and the composition for at least 10 examples, showing clear intermediate variable peaks.

Figure 4: Aggregate processing signatures for tasks with fewer than 10 successful compositional examples, showing less pronounced intermediate variable signals.

Figure 5: Processing signatures for cases where all hops are solved but the composition is not, indicating absence or weakness of intermediate variable representations.

Embedding Space Linearity and Mechanism Selection

A key finding is the strong correlation between the geometry of the embedding space and the processing mechanism employed. When a linear mapping exists from $x$ (input embedding) to $g(f(x))$ (output unembedding), the model tends to favor the direct mechanism. Linearity is quantified via least squares regression and cosine similarity reconstruction accuracy.

Figure 6: Correlation ($r^2$) between embedding space linearity and presence of intermediate variables across tasks.

Tasks with high linearity (e.g., movie-director-birthyear) predominantly exhibit direct processing, while those with low linearity (e.g., antonym-spanish) favor compositional processing. The distribution of compositionality metrics is bimodal, indicating that most examples within a task are processed consistently by one mechanism.

Causal Interventions and Mechanism Validation

Activation patching experiments demonstrate that intermediate variable representations have causal effects on model predictions, particularly in tasks with compositional processing signatures. Patching activations from one task into another alters the predicted outputs in a manner consistent with the presence of the intermediate variable, validating the interpretability findings.

Figure 7: Causal effects on predicted values after patching compositional processing signatures.

Figure 8: Causal effects on predicted values after patching direct processing signatures.

Scaling and Model Architecture Effects

Increasing model size (layers and parameters) yields diminishing improvements in compositionality gap reduction. The gap persists even in instruction-tuned and reasoning-augmented models, suggesting that architectural scaling alone is insufficient for robust compositional generalization.

Figure 9: Diminishing improvements to compositionality gap with increased model size for OLMo 2 and Llama 3 families.

Theoretical and Practical Implications

The results challenge the assumption that compositional behavior in LLMs necessarily arises from compositional mechanisms. The observed duality—models sometimes compute intermediate variables and sometimes do not—suggests that compositionality in behavior does not entail compositionality in representation or process. The choice of mechanism is mediated by the representational structure induced by pretraining, particularly the linearity of relations in embedding space.

This has implications for the design of neuro-symbolic and hybrid systems, as well as for theories of cognitive architecture. The findings support a "bottom-up" approach to understanding compositionality, where primitives are defined by representational adequacy rather than explicit symbolic structure.

Limitations and Future Directions

The paper is limited to single forward-pass computations in Llama 3 (3B) and similar models, with a focus on tasks amenable to current interpretability methods. Further research should extend these analyses to larger models, alternative architectures, and a broader range of compositional functions. The relationship between processing mechanism and compositional generalization remains an open question, warranting causal intervention studies and more granular behavioral analyses.

Conclusion

This paper provides a detailed characterization of how LLMs compose functions, revealing a persistent compositionality gap and a duality of processing mechanisms. The choice between compositional and direct computation is strongly mediated by embedding space linearity, with practical and theoretical consequences for model design and cognitive modeling. The results underscore the need for further investigation into the internal mechanisms of LLMs and their implications for compositional generalization and systematicity.