Discovering a Zeta Map Algorithm on Dyck Paths via Mechanistic Interpretability


Abstract

Machine learning is increasingly used in mathematical discovery, but in mathematics the desired output is often not a prediction itself, but an explicit construction that can be checked independently. We study this setting through the zeta map on Dyck paths, a classical bijection in the combinatorics of the \(q,t\)-Catalan numbers. We train a deliberately small one-layer, one-head encoder–decoder transformer on this map and analyze its learned computation using mechanistic interpretability tools, including decoder cross-attention analysis, linear probing, and causal intervention. The analysis reveals a level-based mechanism: encoder representations make path levels linearly accessible, while the decoder selects and traverses input positions in a structured way. Translating these signals into combinatorics leads to the scaffolding map, an explicit peak-centered traversal algorithm for Dyck paths. We prove that this algorithm agrees with the zeta map, modulo a reversal convention in the labeling. This gives a controlled example of AI-assisted mathematical discovery in which mechanistic interpretability turns model behavior into a precise, human-verifiable combinatorial algorithm.

1 Introduction↩︎

Machine learning is increasingly used in mathematical research [1][3], but many successful applications still treat a model primarily as a high-accuracy predictor or generative model. For mathematical discovery, prediction alone is typically insufficient: the desired output is often a conjecture, construction, or explicit algorithm that a mathematician can verify. This paper studies such a human-AI workflow in a concrete combinatorial setting, the zeta map on Dyck paths. We train transformer models [4] on a known mathematical bijection, interpret the mechanism learned by a deliberately small transformer via mechanistic interpretability tools, and extract from this interpretation a new explicit algorithm.

The zeta map is a central combinatorial object in the theory of \(q,t\)-Catalan numbers because it gives a bijective explanation for the symmetry between the standard Dyck-path statistics appearing in the formulas for \(q,t\)-Catalan numbers. In particular, it converts the \((\mathrm{dinv},\mathrm{area})\) description into the \((\mathrm{area},\mathrm{bounce})\) description, linking two different definitions [5], [6]. Since the zeta map already admits several algorithmic descriptions [5], [7], it provides an ideal test case for AI-assisted mathematical discovery: the data are exactly generable, correctness can be checked exhaustively at moderate sizes, and the desired output is a discrete combinatorial procedure rather than a numerical approximation.

We first show that transformers can learn the zeta map across several experimental settings. We then reduce the architecture to a one-layer, one-head encoder–decoder transformer, for which the learned solution is sufficiently structured to analyze mechanistically. The central quantity revealed by this analysis is the level sequence of the input Dyck path: cross-attention suggests that the decoder selects positions according to level, causal ablations show that many upward-step positions are not directly needed during decoding, and linear probes show that level information is recoverable from encoder representations. These observations led us to formulate the scaffolding map, a peak-centered traversal algorithm for Dyck paths. Unlike the classical area-sequence descriptions of the zeta map, this procedure is organized directly in terms of levels and local traversal. To the best of our knowledge, this is a first example of using mechanistic interpretation of a trained neural sequence model to extract a new explicit bijective description of a well-studied combinatorial map.

A broader goal of this paper is to use the zeta map as a proving ground for AI in combinatorics research, namely to demonstrate the ability to discover combinatorial bijections and algorithms from trained models. While the precise details may differ from one application of AI in math to another, the concepts articulated in this paper are widely applicable within combinatorics research. The model supplies recurring structural signals; the human investigator interprets these signals as combinatorial operations; additional interventions test whether the interpretation is causally relevant; and the resulting procedure becomes a mathematical object that can be stated and proved independently of the neural network.

2 Related Work↩︎

2.0.0.1 AI-Assisted Discovery.

AI methods have shown increasing promise as search tools for mathematical discovery. Reinforcement-learning approaches have been used to find explicit examples and counterexamples in extremal combinatorics, graph theory, and knot theory [1], [8]. More recent work further studies how reinforcement-learning agents navigate difficult mathematical search spaces with sparse high-reward instances [9], [10]. Complementary generative model-guided search and LLM-based coding agents have also been used to discover extremal mathematical structures and improved algorithms [2], [3], [11]. Our work differs in focus: rather than using AI to search for examples, we use mechanistic interpretability to extract an explicit combinatorial algorithm from a trained model.

2.0.0.2 Mechanistic Interpretability.

Mechanistic interpretability seeks to reverse engineer the algorithms implemented by trained neural networks. Relevant precedents in mathematical settings include analyses of small transformers trained on modular arithmetic and group operations [12][14] and arithmetic and enumerative geometry tasks [15], [16]. There are recent efforts on automating discovery such as automated circuit discovery [17], automated neuron or feature explanation [18], [19], and activation-to-language decoding methods [20], [21].

3 Mathematical Background↩︎

3.1 Dyck paths and \(q,t\)-Catalan statistics↩︎

Dyck paths are fundamental objects in combinatorics. A Dyck word of semilength \(n\) is a binary word \(w=(w_i)_{i=1}^{2n}\) of length \(2n\) consisting of exactly \(n\) \(1\)’s such that every initial segment contains at least as many \(1\)’s as \(0\)’s. We write \(\mathrm{Dyck}(n)\) for the set of such words where \(\lvert \mathrm{Dyck}(n)\rvert\) is the well-known \(n\)th Catalan number. Equivalently, by interpreting \(1\) as a North step and \(0\) as an East step, a Dyck path is a lattice path from \((0,0)\) to \((n,n)\) that stays weakly above the diagonal \(y=x\), with an example shown in 1.

Figure 1: An example of Dyck path of semilength 5.

The \(q,t\)-Catalan numbers refine the Catalan numbers and admit two Dyck-path expansions involving the statistics \(\mathrm{area}\), \(\mathrm{bounce}\), and \(\mathrm{dinv}\): \[\begin{align} C_n(q,t) &= \sum_{w\in \mathrm{Dyck}(n)} q^{\mathrm{area}(w)}t^{\mathrm{bounce}(w)} \nonumber \\ &= \sum_{w\in \mathrm{Dyck}(n)} q^{\mathrm{dinv}(w)}t^{\mathrm{area}(w)}. {\label{eq:two-formulas}} \end{align}\tag{1}\] These formulas were discovered independently by [5], [6]. Here \(\mathrm{area}(w)\) counts the cells between the path and the diagonal, \(\mathrm{bounce}(w)\) is computed from Haglund’s bounce path, and \(\mathrm{dinv}(w)\) counts certain diagonal inversions in the area sequence of \(w\).

The equivalence of the two definitions motivates the search for a bijection that exchanges these statistics. Haglund’s zeta map [5] is a bijection \(\zeta\), satisfying \[(\mathrm{dinv}(w),\mathrm{area}(w)) = (\mathrm{area}(\zeta(w)),\mathrm{bounce}(\zeta(w))).\] Although the zeta map is unique, it has several combinatorial interpretations, including descriptions in terms of area sequences [5], [7] and the sweep map [22], [23], which describes its inverse. We do not aim to replace these interpretations. Instead, we ask whether a model trained on the map can reveal an alternative algorithmic organization. We include further background on the \(q,t\)-Catalan numbers and the zeta map in 8.

For a Dyck word \(w=(w_i)_{i=1}^{2n}\), we define its level sequence by \(\ell_0=0\) and \[\ell_{i}=\begin{cases} \ell_{i-1}+1, & \text{ if }w_i=1,\\ \ell_{i-1}-1, & \text{ if }w_i=0. \end{cases}\] Levels later become a central latent combinatorial quantity identified by our interpretability analysis.

4 Experimental Setup↩︎

4.0.0.1 Experiments.

For fixed semilength \(n\), we construct a supervised dataset of Dyck words and their images. More precisely, for each input Dyck word \(w\), the target sequence is the image produced by the SageMath [24] function area_dinv_to_bounce_area_map(). Under the labeling convention used by this function, discussed further in 11, the target agrees with \(\mathrm{rev}\circ\zeta(w)\), where \(\mathrm{rev}\) denotes reading the Dyck word backward and swapping \(0\)’s and \(1\)’s.

Using a fixed architecture with an embedding dimension of \(256\), 4 encoder layers, and 4 decoder layers (each with 8 attention heads), the models achieved \(>99\%\) accuracy in learning the zeta map, that is, predicting the image of a Dyck path under the zeta map for datasets of semilengths \(n = 11, 12, 13, 14, 15, 16\), with sizes \(58{,}786\), \(208{,}012\), \(742{,}900\), \(2{,}674{,}440\), \(9{,}694{,}845\), and \(35{,}357{,}670\), corresponding to the Catalan numbers \(C_n\).

We later varied the architecture and the hyperparameters and explored with a selective subset of

  • encoder-only, decoder-only, and encoder–decoder architectures;

  • the number of transformer layers, ranging from 1 to 4;

  • the number of attention heads, ranging from 1 to 8,

and the model was able to obtain near perfect accuracy for almost all settings.

4.0.0.2 Compression.

We then trained a minimal model with embedding dimension \(128\), one encoder layer, one decoder layer, and one attention head. The model has \(339{,}716\) parameters and achieved perfect or near-perfect accuracy across repeated runs. This model trained on the \(n = 13\) dataset, which we call the Minimal Dyck Transformer, is the focus of the interpretability analysis below. We include architectural details of Minimal Dyck Transformer in 9.

5 Interpreting the Minimal Dyck Transformer↩︎

5.1 Cross-attention patterns↩︎

While several mechanistic inspections reveal useful structure, as shown in 10, the most informative internal signal is decoder cross-attention. We dynamically track, at each generation step, which input positions receive the largest decoder cross-attention weights. An example is shown in 2. During generation, the decoder repeatedly attends to a small and structured subset of input positions. Across examples, we observe three recurring patterns.

Figure 2: Representative decoder cross-attention matrices for the Minimal Dyck Transformer. The recurring sparse patterns indicate that generation is organized by level-based selection rather than by an unstructured lookup table.

5.1.0.1 Highest-level selection.

At early output steps, most prominently observed in the attention visualization when the decoder generates the first step of the output Dyck word, the decoder places attention on input positions whose levels are maximal among relevant positions. An example can be seen in 3 (a) and 3 (b). This suggests that level information has already been computed by the encoder and is being used by the decoder as a selection criterion.

a

b

c

Figure 3: Example of decoder cross-attention and its causal role. Left: input, target, and model-generated Dyck paths, with attended input positions highlighted in yellow. Middle: sparse cross-attention used when generating the first output step on the left. Right: causal ablation results showing negligible decreases in accuracy, measured by comparing the generated output path with the correct target path up to the \(k\)th step, after ablating all cross-attention to input North steps..

5.1.0.2 Systematic avoidance of many upward steps.

Positions with \(w_i=1\) receive no cross-attention throughout decoding, hence the decoder cannot directly “read off’’ level and direction information from upward steps: ignoring \(w_i = 1\) prevents the model from directly accessing the upward steps’ levels and directions. One way for the decoder to compare information of such positions is that the encoder has already aggregated and stored them in the hidden states of the \(w_i = 0\) positions. However, we show this is not the case in 5.2. Instead, it appears to rely heavily on selected positions and reconstructs the rest through a structured traversal.

5.1.0.3 Triangular sequential shift.

As decoding proceeds, many examples exhibit the triangular attention pattern shown in 2. This pattern suggests that, after selecting a high-level position \(i\) with East-step (\(w_i=0\)), the model often shifts attention to the next eligible position \(i+1\) also with East-step. By the Dyck path structure, this position satisfies \(\ell_{i+1}=\ell_i-1\), indicating that the model moves rightward through nearby positions with sequentially decreasing levels.

5.2 Ablation and probing↩︎

We next test whether these visual patterns correspond to functionally relevant structure.

5.2.0.1 Causal ablation.

We mask cross-attention to input positions with \(w_i=1\) during decoding. This causal intervention causes only a negligible degradation in accuracy shown in 3 (c), supporting the hypothesis that many upward-step positions are not essential carriers of decoding-time information.

5.2.0.2 Linear probing.

We train linear probes [25] from encoder hidden states to the corresponding level values \(\ell_i\). The probes recover level information with high predictive accuracy of 84%, indicating that the encoder representation makes this combinatorial statistic linearly accessible. This supports a two-stage interpretation: the encoder computes level-like information, and the decoder uses it to select and traverse positions.

These observations suggest that the model is not memorizing the zeta map as an arbitrary lookup table. Instead, it appears to organize generation by levels, starting from peaks and then traversing nearby positions. Abstracting this behavior led us to the following explicit procedure, which we call the scaffolding map.

6 The Scaffolding Map↩︎

We now state the algorithm suggested by the Minimal Dyck Transformer. Let \(w=(w_i)_{i=1}^{2n}\) be a Dyck word, and compute levels \(\ell_i\) as above. Let \[R=\{i\in\{1,\ldots,2n\}: w_i=0\}\] be the set of East-step positions. A peak position is an index \(i\) such that \((w_i,w_{i+1})=(1,0)\).

Informally, the algorithm starts from peaks at high levels and releases two types of agents that move along adjacent East-step and non-East-step positions. Reading the symbols encountered by these agents, level by level from top to bottom, produces the output word.

6.0.0.1 Scaffolding map.

  1. Group peak positions by their level \(\ell_i\).

  2. Initialize \(\textit{out}=[\,]\), \(\textit{agents}=[\,]\), and set the current level to the maximum level of any peak.

  3. While \(|\textit{out}|<2n\):

    1. Let \(\textit{queue}\) be the union of the current agents and the peak positions at the current level.

    2. Sort \(\textit{queue}\) in decreasing order. For each \(i\) in the queue in this order, append \(w_i\) to \(\textit{out}\).

    3. Update all existing agents simultaneously. If an agent is at a position \(i\) in \(R\), move it to \(i+1\) when this is a valid position in \(R\); otherwise remove it. If an agent is at a position \(i\) not in \(R\), move it to \(i-1\) when this is a valid position not in \(R\); otherwise remove it.

    4. For each peak \(j\) at the current level, spawn new agents at \(j+1\) if \(j+1\in R\) and at \(j-1\) if \(j-1\notin R\), when valid.

    5. Decrease the current level by \(1\).

  4. Return \(\textit{out}\).

We call this the scaffolding map because the upward steps build the level structure, while the output is generated by agents traversing this scaffold from peak positions. The name emphasizes the separation between the latent support structure and the subsequent traversal. Although the transformer is trained at a fixed semilength \(n\), the resulting scaffolding map is defined uniformly for Dyck paths of all semilengths.

Theorem 1. The scaffolding map is equivalent to the zeta map up to a reversal convention in the labeling. Precisely, denote the scaffolding map by \(\varphi\). Define \(\mathrm{rev}\) to be flipping the Dyck path with respect to \(y= -x\), or reading the Dyck word backward with \(0\) and \(1\) swapped. Then \[\varphi = \mathrm{rev} \circ \zeta .\]

The full proof is provided in 11. The proof strategy is to compare the scaffolding map with the area-sequence description of \(\mathrm{rev}\circ\zeta\) level by level. The model-derived algorithm is therefore not presented as a replacement for proof, but as a source of a precise theorem whose verification can be completed by standard combinatorial arguments.

7 Discussions↩︎

7.0.0.1 Summary.

This work uses the zeta map on Dyck paths as a controlled test case for AI-assisted mathematical discovery. By compressing the model to a minimal one-head encoder–decoder transformer, we obtain a learned solution whose internal structure can be inspected directly. The resulting attention, causal ablation, and probing evidence point to a level-based computation, which we translate into the scaffolding map: an explicit peak-centered traversal algorithm for Dyck paths. This illustrates a workflow in which a neural model is not the final mathematical object, but a source of structured evidence that guides the formulation of a new theorem.

7.0.0.2 Limitations.

The present study has several limitations. First, attention patterns alone are not proof of mechanism [26], [27]; this is why we pair them with causal ablations, probes, and an independent combinatorial formulation. Second, the transformer we interpret appears to implement a “fuzzy” version of the scaffolding algorithm. This is reflected, for example, in the imperfect level-probing results. The learned mechanism may become cleaner with longer training, as observed in related settings [12]. More importantly, the scaffolding map was extracted from a particularly simple trained model. Larger models may implement complicated algorithms, and we also observe even small models may admit multiple internal solutions across random seeds [14], [28].

7.0.0.3 Future Directions.

Future work could automate more of this discovery loop and adapt it to larger, less transparent models. Automated circuit-discovery methods and interpretability agents [17], [19], [29], together with activation-to-language methods [20], [21], could help identify recurring internal mechanisms and translate them into candidate symbolic rules. Coupling these tools with theorem provers may further support the transition from model-derived hypotheses to formally verified combinatorial algorithms.

8 \(q,t\)-Catalan Numbers and the Zeta Map↩︎

8.1 \(q,t\)-Catalan Numbers↩︎

We start this section by defining the main objects of interest, Dyck paths.

Definition 1. A Dyck path of semilength \(n\) is a length \(2n\) binary sequence \(w = (w_i)_{i=1,\ldots,2n}\) with the property that each initial segment \((w_i)_{i=1,\ldots,k; \;k \leq 2n}\) of \(w\) contains at least as many \(1\)’s as \(0\)’s. We denote the set of all Dyck paths of semilength \(n\) by \(Dyck(n)\).

Remark 2. Equivalently, a Dyck path of semilength \(n\) is a lattice path from \((0,0)\) to \((n,n)\) consisting of North \((0,1)\) and East \((1,0)\) steps that stays on or above the line \(y=x\). Under this equivalence, a \(1\) in the binary representation corresponds to a North step, and a \(0\) corresponds to an East step. We will use \(w\) to refer to the binary sequence and lattice path representation of a Dyck path interchangeably.

It is well known that the number of Dyck paths of semilength \(n\) is the \(n\)th Catalan number \[C_n = \frac{1}{n+1} {2n \choose n}.\]

There is another equivalent formalization of Dyck paths in terms of the number of complete boxes between the lattice path representation of a Dyck path and the diagonal \(y=x\) in each row, from bottom to top, which is convenient for calculating the statistics in Theorem 1.

Definition 2. The area sequence of a Dyck path \(w\) of semilength \(n\) is the sequence \((a_i)_{i=1,\ldots,n}\) where \(a_i\) counts the number of complete boxes in row \(i\) between the Dyck path and the diagonal \(y=x\), where the rows are enumerated from bottom to top.

Every Dyck path \(w\) uniquely determines an area sequence \((a_i)\) and vice versa. An example of the lattice path, binary sequence, and area sequence representations of a Dyck path is given in Example 1.

We now define the three statistics which appear in 1 .

Definition 3. Let \(w = (w_i)\) be a Dyck path of semilength \(n\) and \((a_i)\) be its area sequence. The area of \(w\) is the nonnegative integer \[area(w) = \sum_{i=1}^{n} a_i.\] The dinv (or diagonal inversion) of \(w\) is the nonnegative integer \[dinv(w) = \#\{ (i,j) \;| \;i < j, a_i = a_j\} + \#\{ (i,j) \;| \;i < j, a_i = a_j + 1 \}.\]

Definition 4. Let \(w = (w_i)\) be a Dyck path of semilength \(n\). The bounce of \(w\) is the nonnegative integer calculated in the following way:

  1. Shoot a billiard from \((0,0)\) heading North

  2. When you hit the start of an East step, turn East and travel until you hit the diagonal

  3. Turn North and repeat until you reach \((n,n)\).

  4. Record the places you hit the diagonal \((0,0), (j_1,j_1),\ldots, (j_k,j_k), (n,n)\).

Then \[bounce(w) = \sum_{i=1}^k (n - j_i).\]

Definition 5. The \(n^{th}\) \(q,t\)-Catalan number is \[C_n(q,t) = \sum_{w \in Dyck(n)} q^{area(w)}t^{bounce(w)} = \sum_{w \in Dyck(n)} q^{dinv(w)}t^{area(w)} \in \mathbb{Z}_{>0}[q,t].\]

8.2 Haglund’s Zeta Map↩︎

The zeta map uses the area sequence/bounce path [5], [7], and can be described by the following algorithm:

  1. Compute the area sequence \(a_1, a_2, \ldots, a_n\) of \(w\) (row lengths from bottom).

  2. Set \(b = \text{max}\{ a_i \}+1\).

  3. For each \(k=0, 1, \ldots , b\), scan the sequence \(a_1, a_2, \ldots, a_n\) from left to right and record:

    1. a \(0\) for each occurrence of \(k-1\),

    2. a \(1\) for each occurrence of \(k\).

    Since the area sequence contains no entries equal to \(-1\) or \(b\), this procedure starts with a \(1\) when \(k=0\) and ends with a \(0\) when \(k=b\).

Below is an example of the image of the zeta map of a Dyck path of semilength 8.

Example 3. Consider the Dyck path \(w\) below. \[\begin{figure}\includegraphics[width=0.8\textwidth]{_pdflatex/pmangrwk.png}\label{shknzeay}\end{figure}\qquad{(1)}\] The area sequence is \((0,1,2,2,2,3,1,2)\). Scanning from left to right for \(-1\) and \(0\), we find no occurrence of \(-1\), and the single occurrence of \(0\) produces \(1\) in the image under the zeta map. Next, scanning for \(0\) and \(1\), we obtain the subsequence \(0,1,1\), which yields \(0,1,1\) in the image. Scanning for \(1\) and \(2\) gives \(1,2,2,2,1,2\), producing \(0,1,1,1,0,1\). Similarly, scanning for \(2\) and \(3\) yields \(2,2,2,3,2\), which produces \(0,0,0,1,0\). Finally, scanning for \(3\) and \(4\) gives \(3\), producing \(0\).

\[\begin{figure}\includegraphics[width=0.8\textwidth]{_pdflatex/ngzdrhkx.png}\label{iwarlzox}\end{figure}\qquad{(2)}\]

9 Architecture of Minimal Dyck Transformer↩︎

We use a deliberately small encoder–decoder transformer. The goal is to keep architectural capacity minimal while still permitting explicit inspection of attention patterns.

9.0.0.1 Model specification.

The model is a standard encoder–decoder transformer without dropout:

  • Embedding/model width: \(d_{\text{model}} = 128\).

  • Feed-forward width: \(d_{\text{ff}} = 256\) with GELU nonlinearity.

  • Layers: \(1\) encoder block and \(1\) decoder block.

  • Attention heads: \(1\).

  • Dropout: \(0\).

  • Maximum length: positional embeddings defined up to \(\texttt{max\_len}=128\).

It has approximately \(3.66\times 10^5\) trainable parameters.

9.0.0.2 Token and positional embeddings.

For a token sequence \((x_1,\dots,x_L)\), the encoder input is \[\mathbf{H} = \mathrm{Embed}_{src}(x_{1:L}) + \mathrm{PosEmbed}(1{:}L),\] and the decoder input is defined similarly, using a distinct target embedding table \(\mathrm{Embed}_{tgt}\). Positional encodings are learned, i.e., given by an embedding lookup indexed by position; sinusoidal encoding is disabled in this configuration.

9.0.0.3 Attention mechanism.

All attention sublayers use scaled dot-product attention. For queries \(Q\), keys \(K\), and values \(V\), \[\mathrm{Attn}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_{\text{head}}}} \right)V.\] Decoder self-attention is causal: when no mask is explicitly provided, an upper-triangular causal mask is applied to prevent attention to future target positions. Encoder self-attention and decoder cross-attention are non-causal.

9.0.0.4 Transformer blocks (post-norm).

Each encoder block applies self-attention and a position-wise feed-forward network with residual connections and post-layer normalization: \[\mathbf{H} \leftarrow \mathrm{LN}\!\big(\mathbf{H} + \mathrm{SelfAttn}(\mathbf{H})\big), \quad \mathbf{H} \leftarrow \mathrm{LN}\!\big(\mathbf{H} + \mathrm{FFN}(\mathbf{H})\big).\] Each decoder block applies (i) causal self-attention, (ii) encoder–decoder cross-attention, and (iii) the same FFN, each followed by residual addition and post-norm: \[\mathbf{Z} \leftarrow \mathrm{LN}\!\big(\mathbf{Z} + \mathrm{SelfAttn}_{\text{causal}}(\mathbf{Z})\big),\; \mathbf{Z} \leftarrow \mathrm{LN}\!\big(\mathbf{Z} + \mathrm{CrossAttn}(\mathbf{Z}, \mathbf{H}_{enc})\big),\; \mathbf{Z} \leftarrow \mathrm{LN}\!\big(\mathbf{Z} + \mathrm{FFN}(\mathbf{Z})\big).\] The FFN is \[\mathrm{FFN}(\mathbf{x}) = W_2\,\mathrm{GELU}(W_1\mathbf{x}+b_1)+b_2,\] with \(W_1\in\mathbb{R}^{128\times256}\) and \(W_2\in\mathbb{R}^{256\times128}\).

9.0.0.5 Output head.

Decoder hidden states are mapped to next-token logits by a linear projection \(\mathbb{R}^{128}\rightarrow\mathbb{R}^{4}\) applied at each position. No weight tying is used between the embeddings and the output projection.

10 Other Embedding And Attention Visualizations↩︎

In this section, we present some visualizations that were helpful when investigating the internal mechanism of the transformer.

10.0.0.1 Additional Cross-Attention Examples.

We include additional examples illustrating the three recurring attention patterns, showing that they are common phenomena rather than hand-picked examples: highest-level selection, systematic avoidance of many positions with \(w_i=1\), and triangular sequential shift.

Figure 4: Cross-attention matrices at a fixed output step 23 generation for a batch of random samples, illustrating that the triangular shift pattern occurs consistently across examples.
Figure 5: Level-attention comparison across examples. Decoder cross-attention at the first output step consistently concentrates on high-level input positions.

10.0.0.2 Encoder self-attention.

Beyond decoder cross-attention, we also inspect encoder self-attention to understand how the model organizes input positions before decoding. The resulting attention maps show recurring block-like structure partly due to the binary nature of Dyck paths.

Figure 6: Encoder self-attention matrices across a batch of examples. The recurring block-like structure reflects the binary organization of Dyck words and suggests that the encoder forms structured position-level representations.

10.0.0.3 Learned positional embeddings.

We further visualize the learned encoder and decoder positional embeddings using PCA. These plots show that the model learns nontrivial positional geometry: for the encoder embeddings, the first principal component separates early positions from the main cluster, while the second principal component largely follows the ordering of input positions. This ordered geometry may help explain the block-like structure observed in the encoder attention maps. Since nearby positions have similar positional embeddings, their position-dependent contributions to attention are similar; consequently, much of the remaining variation in the attention pattern is driven by the binary token value, \(0\) or \(1\). This combination of smooth positional variation and token-level separation naturally produces block-like attention structure.

Figure 7: PCA of the learned encoder positional embeddings. The second principal component largely follows the ordering of input positions, while the first component separates some early positions from the main cluster.
Figure 8: PCA of the learned decoder positional embeddings. Decoder positions form a more regular spiral trajectory in the first two principal components [4], [30].

11 Proof of the Scaffolding–Zeta Equivalence↩︎

To begin with, we clarify conventions and fix notations.

Our scaffolding map, denoted by \(\varphi\), coincides with area_dinv_to_bounce_area_map() in SageMath, which we used to generate our datasets. However, it does not agree with Haglund’s zeta map, denoted by \(\zeta\), due to a difference in conventions: the bounce path is defined in the reverse direction in SageMath compared to Haglund’s formulation. In this section, we prove \[\varphi = \mathrm{rev} \circ \zeta ,\] where \(\mathrm{rev}\) is reading the Dyck word backward with \(0\) and \(1\) swapped.

Let \(a(w)= (a_1, a_2, \dots , a_n)\) be the area sequence of a Dyck path \(w\), and set \(b = \text{max}\{ a_i \}+1\). Recall the algorithm for \(\zeta\) described at the beginning of Section 8.2.

For each \(k=0, 1, \ldots , b\), scan the sequence \(a_1, a_2, \ldots, a_n\) from left to right and record:

  1. a \(0\) for each occurrence of \(k-1\),

  2. a \(1\) for each occurrence of \(k\).

Accordingly, an algorithm for \(\mathrm{rev} \circ \zeta\) is given by the following.

For each \(k=b, b-1, \ldots , 0\), scan the sequence \(a_1, a_2, \ldots, a_n\) from right to left and record:

  1. a \(0\) for each occurrence of \(k\),

  2. a \(1\) for each occurrence of \(k-1\).

From the algorithm of \(\mathrm{rev} \circ \zeta\), for each \(k=b, b-1, \ldots , 0\), we obtain a subsequence \(P_k\) of \(\tilde{a}(w)\) consisting of \(k\) and \(k-1\), where \(\tilde{a}(w) = (a_n, a_{n-1}, \dots , a_1)\). (Recall that we scan \(a(w)\) from right to left in the algorithm of \(\mathrm{rev} \circ \zeta\).) For example, the Dyck path in Example 3 (see 9 below) yields \[P_4= (3), \quad P_3= (2,3,2,2,2), \quad P_2= (2,1,2,2,2,1), \quad P_1 = (1,1,0), \quad P_0 =(0).\] These subsequences completely determine the image of \(\mathrm{rev} \circ \zeta\).

On the other hand, in the algorithm of the scaffolding map \(\varphi\) described in Section 6, each \(w_i\) of \(w=(w_i)\) is associated with level \(\ell_i\), \(1 \le i \le 2n\). Define \[b_i = \begin{cases} \ell_i -1 & \text{if } w_i =1, \\ \ell_i & \text{if } w_i=0 . \end{cases}\] For each \(1 \le i \le 2n\), we associate position \(i\) with \(a_r\) for some \(r\) in \(a(w) = (a_1, \dots , a_n)\) in the following way; in particular, we will have \(b_i =a_r\).

If \(w_i=1\) and position \(i\) belongs to the \(r^{\mathrm{th}}\) row, then we associate position \(i\) with \(a_r\) and have \(b_i= \ell_i-1 = a_r\). If \(w_i=0\), then we slide with slope \(1\) down to the south-west direction until we hit a position, say \(j\), with \(w_j=1\). Suppose that position \(j\) belongs to the \(r^{\mathrm{th}}\) row. Then we associate position \(i\) with \(a_r\) and obtain \(b_i = \ell_i=a_r\) from the symmetry of rows and columns, noting that \(b_i\) is equal to the length of the column that contains position \(i\). For example, in Example 3, position 10, marked with \(\times\), is associated with \(a_5=2\), and position 11, marked with , is associated with \(a_2=1\).

Figure 9: Dyck path with marked positions used in the proof.

In step (3)(a) of the scaffolding algorithm, the queues \(Q_b, Q_{b-1}, \dots, Q_0\) are formed, where \(b =\max \{\ell_i \}= \max\{a_i\} +1\) as before. For \(k=b,b-1, \dots , 0\), each \(Q_k\) consists of positions \(i_1, i_2, \dots\). Replace them with the corresponding \(b_{i_1}, b_{i_2}, \dots\), and denote the resulting sequence by \(Q'_k\). For example, in Example 3, we have \[\begin{align} & Q_4= (8), \quad Q_3= (13,9,7, 5,3), \quad Q_2= (14, 12, 10, 6, 4,2), \quad Q_1 = (15, 11, 1), \quad Q_0 =(16), \\ & Q'_4= (3), \quad Q'_3= (2,3,2,2,2), \quad Q'_2= (2,1,2,2,2,1), \quad Q'_1 = (1,1,0), \quad Q'_0 =(0) . \end{align}\] The sequences \(Q_k\), or the sequences \(Q'_k\) determine the image of the scaffolding map \(\phi\).

We now claim that \(P_k = Q'_k\) for \(k = b, b-1, \dots, 0\). Indeed, in each \(P_k\), the occurrences of \(k-1\) correspond to entries \(b_i = k-1\) in \(Q'_k\), arising from peaks or positions \(i \notin R\) in \(Q_k\); the occurrences of \(k\) correspond to positions \(i \in R\), where \(b_i = k = a_r\), and the \(r^{\mathrm{th}}\) row is associated with \(i\) as described above in the penultimate paragraph.

Since \(P_k\) and \(Q'_k\) determine the images of \(\mathrm{rev}\circ \zeta\) and \(\phi\), respectively, the claim implies that \(\mathrm{rev} \circ \zeta = \phi\), as desired.

References↩︎

[1]
Gukov, S., Halverson, J., Ruehle, F., and Sułkowski, P. Learning to unknot. Machine Learning: Science and Technology, 2 (2): 025035, 2021.
[2]
Charton, F., Ellenberg, J., Wagner, A., and Williamson, G. , 2024.
[3]
Novikov, A., Vũ, N., Eisenberger, M., Dupont, E., Huang, P.-S., Wagner, A. Z., Shirobokov, S., Kozlovskii, B., Ruiz, F. J., Mehrabian, A., et al. Alphaevolve: A coding agent for scientific and algorithmic discovery. arXiv preprint arXiv:2506.13131, 2025.
[4]
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[5]
Haglund, J. Conjectured statistics for the q,t-Catalan numbers. Advances in Mathematics, 175 (2): 319–334, May 2003. ISSN 0001-8708. . URL https://www.sciencedirect.com/science/article/pii/S0001870802000610.
[6]
Haiman, M. . Preprint, University of California, Berkeley, 2000.
[7]
Andrews, G., Krattenthaler, C., Orsina, L., and Papi, P. ad-Nilpotent b-Ideals in sl(n) Having a FixedClass of Nilpotence: Combinatorics and Enumeration. Transactions of the American Mathematical Society, 354 (10): 3835–3853, 2002. ISSN 0002-9947. URL https://www.jstor.org/stable/3072985. Publisher: American Mathematical Society.
[8]
Wagner, A. Z. Constructions in combinatorics via neural networks, 2021. URL https://arxiv.org/abs/2104.14516.
[9]
Butbaia, G., Orland, P., Huang, X., Passaro, D., Fagan, L., Tarquini, M., Dao, H., Eisenbud, D., Shehper, A., and Gukov, S. Hierarchical reinforcement learning for sparse-reward search in commutative algebra. In Forty-third International Conference on Machine Learning, 2026. URL https://openreview.net/forum?id=DF6jVG4fG8.
[10]
Fagan, L., Tarquini, M., Shehper, A., Manko, M., Huang, X., Gruen, A., Butbaia, G., Passaro, D., and Gukov, S. The two-hump problem: Bridging the difficulty gap in mathematical reinforcement learning. In Forty-third International Conference on Machine Learning, 2026. URL https://openreview.net/forum?id=PUQ0rrmuDM.
[11]
Bérczi, G., Hashemi, B., and Klüver, J. Flow-based extremal mathematical structure discovery. arXiv preprint arXiv:2601.18005, 2026.
[12]
Nanda, N., Chan, L., Lieberum, T., Smith, J., and Steinhardt, J. Progress measures for grokking via mechanistic interpretability, 2023.
[13]
Chughtai, B., Chan, L., and Nanda, N. A toy model of universality: Reverse engineering how networks learn group operations. In International Conference on Machine Learning, pp. 6243–6267. PMLR, 2023.
[14]
Zhong, Z., Liu, Z., Tegmark, M., and Andreas, J. The clock and the pizza: Two stories in mechanistic explanation of neural networks, 2023.
[15]
Babei, A., Charton, F., Costa, E., Huang, X., Lee, K.-H., Lowry-Duda, D., Narayanan, A., and Pozdnyakov, A. Learning euler factors of elliptic curves, 2025.
[16]
Hashemi, B., Corominas, R., and Giacchetto, A. Can transformers do enumerative geometry? In International Conference on Learning Representations, volume 2025, pp. 70047–70067, 2025.
[17]
Conmy, A., Mavor-Parker, A., Lynch, A., Heimersheim, S., and Garriga-Alonso, A. Towards automated circuit discovery for mechanistic interpretability. Advances in Neural Information Processing Systems, 36: 16318–16352, 2023.
[18]
Bills, S., Cammarata, N., Mossing, D., Tillman, H., Gao, L., Goh, G., Sutskever, I., Leike, J., Wu, J., and Saunders, W. Language models can explain neurons in language models. https://openaipublic.blob.core.windows.net/neuron-explainer/paper/index.html, 2023.
[19]
Shaham, T. R., Schwettmann, S., Wang, F., Rajaram, A., Hernandez, E., Andreas, J., and Torralba, A. A multimodal automated interpretability agent. In Forty-first International Conference on Machine Learning, 2024.
[20]
Pan, A., Chen, L., and Steinhardt, J. Latentqa: Teaching llms to decode activations into natural language, 2026. URL https://arxiv.org/abs/2412.08686.
[21]
Karvonen, A., Chua, J., Dumas, C., Fraser-Taliente, K., Kantamneni, S., Minder, J., Ong, E., Sharma, A. S., Wen, D., Evans, O., and Marks, S. Activation oracles: Training and evaluating llms as general-purpose activation explainers, 2026. URL https://arxiv.org/abs/2512.15674.
[22]
Armstrong, D., Loehr, N., and Warrington, G. . Advances in Mathematics, 284: 159–185, 2015.
[23]
Thomas, H. and Williams, N. Sweeping up zeta. Selecta Mathematica, 24 (3): 2003–2034, 2018.
[24]
The Sage Developers. SageMath, the Sage Mathematics Software System (Version 9.5), 2026. .
[25]
Alain, G. and Bengio, Y. Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.01644, 2016.
[26]
Jain, S. and Wallace, B. C. Attention is not explanation, 2019. URL https://arxiv.org/abs/1902.10186.
[27]
Wiegreffe, S. and Pinter, Y. Attention is not not explanation, 2019. URL https://arxiv.org/abs/1908.04626.
[28]
Wen, K., Li, Y., Liu, B., and Risteski, A. Transformers are uninterpretable with myopic methods: a case study with bounded dyck grammars, 2023. URL https://arxiv.org/abs/2312.01429.
[29]
Friedman, D., Wettig, A., and Chen, D. Learning transformer programs. Advances in Neural Information Processing Systems, 36: 49044–49067, 2023.
[30]
Wennberg, U. and Henter, G. Learned transformer position embeddings have a low-dimensional structure. In Proceedings of the 9th workshop on representation learning for NLP (RepL4NLP-2024), pp. 237–244, 2024.