How Do Transformers Learn Variable Binding in Symbolic Programs?

Thank you for your feedback that our paper "proposes an interesting paradigm in analyzing the symbolic learning ability of Transformer." We appreciate your critical assessment and have conducted new experiments to address your concerns.

Addressing key concerns:

1. "The model does not distinguish between symbols and values"

We conducted new subspace analysis experiments specifically to address this concern:

1. We applied PCA to the Transformer's residual stream vectors for the right-hand-side tokens of program lines across many different inputs.
2. We trained linear classifiers (with L1 regularization) identify principal components of the residual stream correlated with numerical constants (10 components) and variable names (26 components), respectively.
3. We performed interchange interventions on these selected components, using counterfactual examples that target the numerical constant or the variable name, achieving high success rates for both numerical constants (92.17%) and variable names (87.08%).

New findings: The 2D UMAP projections of these subspaces across three training checkpoints (steps 1200, 17600, 105400) reveal a clear evolution toward separated clusters for these token types. By the final checkpoint, the model distinctly represents symbols and values in different subspaces. See the new the variable-name subspace scatter plot and the numerical-constant subspace scatter plot for more details (direct links to new plots: https://i.imgur.com/ugP0baG.jpeg, https://i.imgur.com/479j8Ch.jpeg).

Action: We will add these new experiments and plots.

2. Questions about training

How are the three phases decomposed?

The three phases emerge naturally during training without external parameter changes. They represent qualitative shifts in strategy that correspond to sharp transitions in performance (from ~12% to ~56% to >99% accuracy). We use a linear decay learning rate schedule with 750 warm-up steps, but the phase transitions do not align with learning rate schedule changes.

What objective is used for training the Transformer?

We use standard autoregressive next-token prediction (causal language modeling) as stated in the Appendix section "Training Setup."

Action: We will add this information more explicitly in the main text.

How to explain Fig. 2c? Why 1-hop accuracy converged slower than those with more hops?

This results from our task structure. In 1-hop programs, the answer is always a numerical constant directly assigned to the queried variable. The model's early "line-1/line-2 heuristic" works well for multi-hop programs where the first line often contains the root value, but performs poorly on 1-hop programs where the answer could be on any line.

Action: We will explain this in the figure caption.

3. Task design

While we acknowledge the controlled nature of our task, its simplicity is intentional. By using a minimal design, we can perform detailed causal analysis that would be challenging with more complex symbolic tasks. Our new generalization experiments show the model learns a truly general mechanism rather than simply memorizing patterns.

We would like to clarify an important aspect of our experimental design: our model is trained entirely from scratch without any pre-training on natural language or programming tasks. When the model first encounters our variable assignment syntax (e.g., "a=1"), it has no prior knowledge of what these symbols mean or how they relate to each other. From this perspective, using alternative symbolic relationships like "a is left of b" would be functionally equivalent - the model would still need to learn these relationships from scratch.

Our focus was on creating a minimal yet challenging test bed where we can precisely analyze how variable binding mechanisms emerge during training. The apparent simplicity of our task masks its fundamental complexity: the model must still develop sophisticated capabilities to track multi-step variable dependencies while ignoring distractors, regardless of the specific symbolic relationship used.

The controlled design was crucial for our mechanistic interpretability goals, allowing us to isolate and analyze the precise causal pathways through which variable binding is implemented in the network.

Additional improvements:

• We will fix all typos, including the dataset split percentages
• We will improve the clarity of Fig. 2d
• We will add clearer explanations of the training objective and methodology

Action: We will implement all these improvements in the camera-ready version.

Thanks again for your helpful feedback! Do these clarifications and planned changes address your concerns?

Thank you for your assessment that we "carefully designed the experiments to empirically validate [our] claims" and for your thoughtful suggestions. Your questions about generalization inspired several new experiments that strengthen our findings.

Does the model actually learn a systematic mechanism?

We appreciate the critical perspective on whether the model indeed learns a systematic mechanism in addition to a shallow heuristic, and conduct generalization experiments to further investigate the matter. Since our model is trained from scratch rather than pre-trained, testing on completely unseen tokens would mean testing on untrained embeddings, which wouldn't be meaningful. Instead, we conducted three rigorous generalization experiments:

• Held-out variable/number combinations: We trained a model where specific variable/number combinations, randomly sampled and constituting ~10% of the original training data, were excluded from training but included in testing. The model successfully generalized to these unseen combinations, demonstrating compositional generalization rather than memorization. See the compositional generalization plot for further details. (direct link to new plot: https://i.imgur.com/esnjS3O.png)

• Program length generalization: Testing on programs of 2-25 lines of assignment (vs. 16 in training), we found excellent generalization on programs where the answer is not on the first line, however when the answer is on the first or second line and the first line heuristic is used we see poor performance on shorter and longer programs. See the program length generalization plot for further details (direct link to new plot: https://i.imgur.com/dIaCWx8.png)

• Chain depth generalization: Testing on programs with 1-13 hops (vs. max 4 in training), we again observed near-perfect generalization, even for 13-hop programs despite our model having only 12 layers. We again see that when the answer is on the first or second line and the first line heuristic the model is worse at higher depth programs. See the chain depth generalization plot for further details. (direct link to new plot: https://i.imgur.com/C2XEVwc.png)

In the paper, we claim that the model uses a shallow heuristic for programs where the answer is in the first or second line, and uses a systematic solution for other programs.

When the answer isn’t on the first or second line, the model generalizes far beyond the structures seen during training – to unseen variable/constant combinations, deeper chains, and longer programs. These results further support our claim that the model learns a "systematic mechanism for dereferencing" rather than memorizing patterns.

Moreover, when the answer is on the first or second line, we see seriously degraded generalization performance. These results further support our claim that these are exactly the inputs that a shallow heuristic is used to solve the task.

Action: We will add these new experiments and plots.

Addressing clarity issues:

• We will provide clearer definitions of technical terminology
• We will add details on how logits are computed at each layer in Figure 3
• We will fix the missing reference in Figure 3(c)
• We will ensure consistent naming and referencing throughout (e.g., "fig. 2a" vs "Fig. 2d")

We will improve the readability of all figures by:

• Increasing font sizes
• Adding clearer labels and legends
• Ensuring consistent formatting

Action: We will implement all these improvements in the camera-ready version.

Thank you for your valuable suggestions that led to these additional experiments! Do these results and planned improvements address your concerns about whether we've demonstrated a truly systematic mechanism?

We thank you for your thorough review and positive feedback that our paper is "straightforward and generally well-written" with experiments that "validate the authors' claims convincingly."

Responses to your questions:

Q1&Q2: Why exactly 17 lines? Did you try testing on longer/shorter programs? What happens with chains longer than 12?

See the response to Reviewer ZyPs for the results of new generalization experiments.

Q3: Did you try reproducing this over multiple training runs?

Yes, we trained additional models with different random seeds. All runs exhibit the same three distinct learning phases with similar transition points, demonstrating the robustness of our findings. See the test set accuracy plot across different training run seeds for more details (direct link to new plot: https://i.imgur.com/cvf0vzO.png).

Action: We will add this new experiment and plot.

Q4: Did you try other distributions?

Yes, we explored different sampling strategies during task design. Our final distribution was deliberately constructed to be as challenging as possible while remaining learnable.

Simpler distributions (e.g., with fewer distractor chains, uniform sampling of variables, or shorter chains) made the task trivially solvable using surface-level heuristics. For instance, without our weighted sampling approach (where we choose chains to extend with probability proportional to chain length cubed), the model could simply learn to associate the answer with the longest variable chain.

Our rejection sampling procedure ensures balance across referential depths while maintaining sufficient distractor chains that branch from the main query chain. This forces the model to genuinely track variable bindings rather than rely on pattern matching.

This challenging distribution was essential for studying true variable binding mechanisms rather than shortcut learning.

Action: We will make this more explicit in the appendix.

Minor points

The logit differences were normalized to a [0,1] scale for visualization clarity. We agree this could be clearer and will explicitly state the normalization in the figure caption.

We acknowledge your concern about figure readability. In the camera-ready version, we will:

• Increase font sizes throughout all figures
• Use more visually distinct colors and line styles in Fig. 2b and Fig. 4
• Add clearer legends and improve visual mapping between colors and values

Thanks again for your helpful feedback! Do these clarifications and planned changes address your concerns?

Thank you for your positive assessment that our paper is "well-written and clear" with claims that are "well-supported." We're especially pleased you highlighted our most interesting finding: that relevant circuits form with "behavior itself being fairly uninteresting."

Response to your question: Linear probing vs. patching

Why do you think linear probing works so poorly yet you are able to see clean effects from patching?

Our linear probing experiment attempts to extract a complete program state from a single vector, while our patching experiments reveal the causal path of specific information through the network. The latter succeeds because the model implements variable binding as a dynamic process of information routing rather than as static state representations.

The poor linear probing results suggest the model doesn't maintain a complete program state in a linearly decodable format. Instead, our patching results reveal that the model learns a more efficient strategy that dynamically tracks only the relevant variable bindings through specialized attention patterns.

Action: We will make this more explicit.

Addressing your other feedback:

• Grammar issues: We will correct all grammatical errors, including:

  • Line 95: "open question of how does" → "open question of how"
  • Line 136: "hand side of the quality" → "equality"
  • Line 291: We'll rephrase the statement about uniform representations

• Figure density: We will redesign Figure 4 with:

  • A larger, clearer legend at the top
  • Thicker, more visually distinct lines
  • Consistent labeling schemes

Action: We will implement all these improvements in the camera-ready version.

Thanks again for your helpful feedback! Do these clarifications and planned changes address your concerns?