Low Rank and Sparse Fourier Structure in Recurrent Networks Trained on Modular Addition

TL;DR

The paper investigates how RNNs learn modular addition and reveals that they implement a Fourier multiplication algorithm, mirroring findings in one-layer transformers. Through Fourier analysis and causal interventions, the authors show the model relies on a sparse set of Fourier frequencies, with six dominant frequencies shared across inputs, hidden states, and outputs, and that key singular vectors align with these frequencies. They demonstrate near-exact Fourier multiplication in the final layer by fitting hidden representations to cosine and sine components at the identified frequencies, with relative errors on the order of $10^{-3}$ to $10^{-2}$. Ablation studies indicate distributed reliance on multiple frequencies, as removing several frequencies degrades performance substantially, while removing a single frequency often leaves functionality intact. These results extend mechanistic interpretability to RNNs and highlight a low-rank, frequency-based computation that could inform understanding of learning dynamics in other algorithmic tasks.

Abstract

Modular addition tasks serve as a useful test bed for observing empirical phenomena in deep learning, including the phenomenon of \emph{grokking}. Prior work has shown that one-layer transformer architectures learn Fourier Multiplication circuits to solve modular addition tasks. In this paper, we show that Recurrent Neural Networks (RNNs) trained on modular addition tasks also use a Fourier Multiplication strategy. We identify low rank structures in the model weights, and attribute model components to specific Fourier frequencies, resulting in a sparse representation in the Fourier space. We also show empirically that the RNN is robust to removing individual frequencies, while the performance degrades drastically as more frequencies are ablated from the model.

Figures (4)

• Figure 1: Unrolled computational graph of the RNN trained to perform modular addition
• Figure 2: Tracking the Fourier coefficients of every node in the computational graph. On the left we plot the coefficients of the embedding and unembedding layers, as well as $h_1$ since it is just a function of $a$. On the right we plot the coefficients of the hidden states $h_2, h_3$ which are functions of $a,b$ (darker colors indicate larger magnitude coefficients). For all tensors, we find that the same $6$ frequencies, $\omega_k = \frac{2\pi k}{P}$ for $k \in \{6, 15, 20, 29, 47, 54\}$ are the only significant coefficients.
• Figure 3: Singular value spectra of model weights are low rank. The number of components $r$ required to maintain $100\%$ model accuracy as follows: embedding matrix $W_E$ : $r=12$ containing $85.4\%$ of the energy , unembedding matrix $W_{fc}$ : $r=12$ containing $95.5\%$ of the energy, input-hidden matrix $W_{ih}$ : $r=12$ containing $91.4\%$ of the energy, and hidden-hidden matrix $W_{hh}$ : $r=32$ containing $89.4\%$ of the energy.
• Figure 4: Identifying the Fourier frequencies associated with different singular vectors $U_{fc}^k, V_{fc}^k$ for $k=1,\ldots,12$ of the unembedding matrix. We find that each consecutive pair of singular vectors is associated with a distinct frequency.