Technical Summary: Discovering a Zeta Map Algorithm on Dyck Paths via Mechanistic Interpretability

Problem Statement

Machine learning is increasingly applied to mathematical research, yet a fundamental gap remains between high-accuracy prediction and the generation of explicit, verifiable mathematical constructions. In combinatorics, the desired output is often a conjecture, construction, or algorithm that can be independently checked, rather than a numerical approximation. This paper addresses this gap by investigating whether mechanistic interpretability tools can extract explicit combinatorial algorithms from trained neural networks.

The specific test case is the zeta map ($\zeta$) on Dyck paths, a classical bijection central to the theory of $q,t$-Catalan numbers. The zeta map exchanges the $(dinv, area)$ statistics for the $(area, bounce)$ statistics. While the zeta map is well-understood and admits several existing algorithmic descriptions (e.g., via area sequences or sweep maps), it serves as an ideal controlled environment for AI-assisted discovery because the data is exactly generable, correctness is exhaustively checkable, and the target is a discrete combinatorial procedure.

Methodology

1. Experimental Setup

The authors trained transformer models to learn the zeta map on Dyck paths of semilength $n$.

• Data: Supervised datasets of Dyck words and their images under the zeta map (specifically $\text{rev} \circ \zeta$, where $\text{rev}$ denotes reading the word backward and swapping 0s and 1s, consistent with the SAGEMATH implementation used).
• Initial Training: Standard encoder-decoder transformers (4 layers, 8 heads) achieved $>99\%$ accuracy across various semilengths ($n=11$ to $16$).
• Compression: To enable mechanistic analysis, the authors trained a Minimal Dyck Transformer: a deliberately small architecture with one encoder layer, one decoder layer, and one attention head (embedding dimension 128, ~340k parameters). This model achieved perfect or near-perfect accuracy.

2. Mechanistic Interpretability Analysis

The core of the methodology involves analyzing the internal workings of the Minimal Dyck Transformer using three primary tools:

• Decoder Cross-Attention Analysis: The authors tracked which input positions received the highest attention weights during the generation of the output sequence.
• Causal Ablation: They masked cross-attention to specific input positions (e.g., upward steps where $w_i=1$) to test if those positions were functionally necessary for decoding.
• Linear Probing: Linear classifiers were trained on encoder hidden states to predict the "level" of each position in the Dyck path, testing whether this combinatorial statistic was linearly accessible in the representations.

3. Algorithm Extraction

Based on the observed patterns in the minimal model, the authors formulated a new explicit algorithm, the Scaffolding Map, and subsequently proved its equivalence to the zeta map.

Key Findings and Contributions

1. Discovery of a Level-Based Mechanism

The interpretability analysis revealed that the model does not treat the zeta map as an arbitrary lookup table but organizes generation via a structured, level-based mechanism:

• Highest-Level Selection: The decoder consistently attends to input positions with maximal levels among relevant positions, particularly at early generation steps.
• Systematic Avoidance of Upward Steps: The decoder rarely attends to positions where $w_i=1$ (North steps). Causal ablation confirmed that masking attention to these positions caused negligible accuracy degradation, implying the decoder does not "read" them directly but relies on aggregated information.
• Triangular Sequential Shift: The attention pattern exhibits a shift from a high-level East-step position ($w_i=0$) to the next eligible East-step position ($i+1$), corresponding to a sequential decrease in level ($\ell_{i+1} = \ell_i - 1$).
• Linear Accessibility: Linear probes confirmed that the encoder representations make the level sequence ($\ell_i$) of the Dyck path linearly accessible (84% predictive accuracy).

2. The Scaffolding Map

Translating these signals into combinatorics, the authors proposed the Scaffolding Map, a peak-centered traversal algorithm:

• Mechanism: The algorithm groups peak positions by their level. It initializes "agents" at these peaks and traverses the path level-by-level from top to bottom.
• Traversal: Agents move along adjacent East-step and non-East-step positions. The output word is generated by reading the symbols encountered by these agents as they traverse the "scaffold" built by the upward steps.
• Structure: The algorithm separates the latent support structure (the levels built by upward steps) from the traversal process that generates the output.

3. Theoretical Verification

The authors provided a formal proof (Theorem 6.1) establishing that the Scaffolding Map is equivalent to the zeta map (modulo the reversal convention used in the dataset).

• Proof Strategy: The proof compares the Scaffolding Map level-by-level with the area-sequence description of $\text{rev} \circ \zeta$. It demonstrates that the subsequences generated by the Scaffolding Map's queueing mechanism ($Q'_k$) exactly match the subsequences generated by the standard zeta map algorithm ($P_k$) when scanning the area sequence in reverse.

Significance and Claims

The paper claims to provide a controlled example of AI-assisted mathematical discovery where mechanistic interpretability transforms model behavior into a precise, human-verifiable combinatorial algorithm.

• Novelty: To the authors' knowledge, this is the first instance of using the mechanistic interpretation of a trained neural sequence model to extract a new explicit bijective description of a well-studied combinatorial map.
• Workflow: The paper articulates a workflow where:
  1. The model supplies recurring structural signals (attention patterns, linear probes).
  2. Human investigators interpret these signals as combinatorial operations.
  3. Interventions (ablations) test causal relevance.
  4. The resulting procedure becomes a mathematical object (the Scaffolding Map) that can be stated and proved independently of the neural network.
• Limitations: The authors are modest about the scope, noting that:
  • Attention patterns alone are not proof; the combination with ablation and formal proof is essential.
  • The learned mechanism in the transformer appears to be a "fuzzy" version of the algorithm (evidenced by imperfect probing results), suggesting that larger or differently trained models might implement more complex or "cleaner" mechanisms.
  • The Scaffolding Map was extracted from a specifically simple model; larger models may implement different internal solutions.

The paper concludes that while the specific details of this discovery may not transfer directly to all mathematical problems, the concept of using mechanistic interpretability to bridge the gap between black-box prediction and explicit combinatorial construction is widely applicable within combinatorics research.