Towards Understanding What State Space Models Learn About Code

State Space Models (SSMs) have emerged as an efficient alternative to the transformer architecture. Recent studies show that SSMs can match or surpass Transformers on code understanding tasks, such as code retrieval, when trained under similar conditions. However, their internal mechanisms remain a black box. We present the first systematic analysis of what SSM-based code models actually learn and perform the first comparative analysis of SSM and Transformer-based code models. Our analysis reveals that SSMs outperform Transformers at capturing code syntax and semantics in pretraining but forgets certain syntactic and semantic relations during fine-tuning on task, especially when the task emphasizes short-range dependencies. To diagnose this, we introduce SSM-Interpret, a frequency-domain framework that exposes a spectral shift toward short-range dependencies during fine-tuning. Guided by these findings, we propose architectural modifications that significantly improve the performance of SSM-based code model, validating that our analysis directly enables better models.

💡 Research Summary

This paper presents the first systematic investigation into what state‑space models (SSMs) actually learn when applied to code‑understanding tasks, and it directly compares these models to the dominant Transformer‑based approaches. The authors focus on encoder‑only architectures—CodeSSM (an SSM‑based model built on the BiGS/S4D backbone) and RoCoder (a Transformer with Rotary Positional Embeddings)—trained under identical pre‑training conditions on large code corpora.

Pre‑training Representation Analysis
Using the classifier‑free DirectProbe methodology, the study probes three relational tasks derived from abstract syntax trees (AST) and data‑flow graphs (DFG): (1) AST distance prediction, (2) AST sibling prediction, and (3) DFG edge prediction. Across all layers, CodeSSM consistently outperforms RoCoder on syntactic metrics (distance and sibling) and matches or exceeds it on semantic edge prediction, especially in deeper layers where long‑range dependencies are crucial. This confirms that SSMs are inherently better at capturing global code structure during unsupervised learning.

Fine‑tuning Dynamics
The models are fine‑tuned on two downstream tasks: Stack‑Overflow question‑answer retrieval (SQA) and type inference. While CodeSSM‑SQA retains its pre‑training advantage and surpasses RoCoder‑SQA, the type‑inference setting reveals a dramatic collapse of CodeSSM’s representations. Layer‑wise probing shows that after fine‑tuning, CodeSSM loses accuracy on short‑range syntactic relations (especially in the early layers), whereas RoCoder actually improves its short‑range performance. This divergence explains why CodeSSM, despite strong pre‑training, underperforms on tasks that require a balanced mix of local and global context.

Spectral Kernel Diagnosis (SSM‑Interpret)
To uncover the root cause, the authors introduce SSM‑Interpret, a novel frequency‑domain analysis of the convolution kernels that drive each SSM layer. By extracting forward and backward kernels from all 12 layers of CodeSSM and applying Fourier transforms, they obtain spectra that reveal the model’s bias toward low‑frequency (long‑range) or high‑frequency (short‑range) token interactions. The key finding is a “spectral shift” during type‑inference fine‑tuning: early‑layer kernels move from low‑frequency dominance to high‑frequency dominance, effectively discarding the long‑range information learned during pre‑training. This shift is absent in RoCoder, whose attention mechanism remains agnostic to frequency and can adapt to the task’s needs.

Architectural Remedies
Guided by the spectral analysis, the paper proposes two concrete modifications to mitigate the shift: (1) frequency‑aware kernel regularization that penalizes excessive high‑frequency energy, and (2) a hybrid routing scheme that combines a short‑window convolution with a global token‑wise aggregation in the early layers. Both variants preserve the low‑frequency component while still allowing the model to capture necessary short‑range patterns. Empirical results show substantial gains: NL‑CodeSearch MRR improves by +5.5, Long‑Context Retrieval MRR by +3.49, and type‑inference F1 by +1.28, thereby validating that the interpretability insights directly translate into better models.

Implications and Future Work
The study demonstrates that SSMs excel at learning global code structure but are vulnerable to task‑specific fine‑tuning that over‑emphasizes short‑range dependencies. Frequency‑domain kernel analysis offers a powerful diagnostic tool for SSMs, complementing the well‑established attention‑map analyses for Transformers. The authors suggest that future research should explore hybrid architectures that leverage SSMs for long‑range context while retaining Transformer‑style mechanisms for precise local reasoning, and should investigate loss‑function designs that maintain a balanced spectral profile throughout fine‑tuning.

In sum, this work bridges a critical interpretability gap, provides a novel analytical framework (SSM‑Interpret), and delivers actionable architectural improvements that advance the state of the art in code‑understanding models.