Protein Circuit Tracing via Cross-layer Transcoders

Abstract: Protein LLMs (pLMs) have emerged as powerful predictors of protein structure and function. However, the computational circuits underlying their predictions remain poorly understood. Recent mechanistic interpretability methods decompose pLM representations into interpretable features, but they treat each layer independently and thus fail to capture cross-layer computation, limiting their ability to approximate the full model. We introduce ProtoMech, a framework for discovering computational circuits in pLMs using cross-layer transcoders that learn sparse latent representations jointly across layers to capture the model's full computational circuitry. Applied to the pLM ESM2, ProtoMech recovers 82-89% of the original performance on protein family classification and function prediction tasks. ProtoMech then identifies compressed circuits that use <1% of the latent space while retaining up to 79% of model accuracy, revealing correspondence with structural and functional motifs, including binding, signaling, and stability. Steering along these circuits enables high-fitness protein design, surpassing baseline methods in more than 70% of cases. These results establish ProtoMech as a principled framework for protein circuit tracing.

• The paper introduces cross-layer transcoders that jointly capture sparse, interpretable latent spaces across transformer layers to serve as a faithful surrogate for pLM computations.
• It achieves 82–89% of original ESM2 performance while identifying minimal, biologically coherent circuits, enabling effective denoising and targeted protein design.
• The results validate ProtoMech’s mechanistic discovery by mapping known biological motifs and advancing both model interpretability and practical applications in protein engineering.

Protein Circuit Tracing via Cross-layer Transcoders: An In-depth Analysis of ProtoMech

Introduction and Motivation

The interpretability of protein LLMs (pLMs), such as ESM2, has become a central concern due to their growing influence in predicting protein structure and function. Despite robust downstream performance, their internal computations—“circuits”—that map sequence input to functional predictions remain obscure. Established approaches, notably sparse autoencoders (SAEs), provide access to interpretable features but fall short of capturing the cross-layer, dynamic computations intrinsic to transformer architectures. The lack of compositional and mechanistic understanding impedes model validation, engineering, and targeted protein design.

ProtoMech addresses this gap by introducing a framework based on cross-layer transcoders (CLTs), which learn sparse, interpretable latent spaces jointly across transformer layers. Unlike standard layer-local approaches (e.g., per-layer transcoders, PLTs), ProtoMech’s CLTs explicitly model dependencies among all preceding layers at each computation step, providing a functionally faithful and highly compressible surrogate of the original model. This strategy enables quantitative circuit tracing, targeted steering, and mechanistic discovery in pLMs, yielding insight into both known and novel biological motifs.

Methodology: Cross-layer Transcoders and Circuit Discovery

ProtoMech is built on the premise that transformer MLP layers perform complex, distributed computation that cannot be disambiguated by independent layerwise surrogates. Instead, the replacement model learns to transcode each MLP as a sparse function of all latent activities in earlier layers (Figure 1).

Figure 1: Overview of ProtoMech. CLTs reconstruct each layer's output from the sparse latents of all preceding layers, capturing the full computational graph.

For ESM2-8M ($L=6$, $d_{model}=320$), ProtoMech uses an encoder-decoder architecture with the following salient details:

• At each layer $\ell$, the pre-MLP hidden state $\mathbf{x}^{\ell}$ is encoded into a $k$-sparse latent vector via a learned projection and TopK selection.
• Each layer's output is decoded as a sum over all earlier layers' latent activations, i.e., $$\hat{\mathbf{y}}^{\ell} = \sum_{\ell'=1}^{\ell} W_{dec}^{\ell' \rightarrow \ell}a^{\ell'} + b_{pre}^{\ell}$$.
• Training minimizes MSE reconstruction error plus an auxiliary loss to prevent dead latents, promoting broad latent utilization and interpretability.

The resulting CLT replacement model closely reconstructs ESM2’s MLP computations. Crucially, the approach supports circuit discovery: ProtoMech employs attribution-based greedy search over latent variables, identifying minimal circuits sufficient to drive task-specific performance, such as protein family classification (F1) or function prediction (Spearman, DMS assays).

Quantitative Performance and Circuit Compression

ProtoMech's replacement model achieves 82–89% of the original ESM2 performance on both family classification and function prediction tasks. For protein family classification, ProtoMech attains F1 = $0.82 \pm 0.19$ (ESM2: $0.92 \pm 0.1$); for function prediction, Spearman = $0.41 \pm 0.19$ (ESM2: $0.50 \pm 0.22$). This substantially outperforms PLT baselines (family F1 $\approx 0.50$, function Spearman $\approx 0.38$), validating the necessity of modeling cross-layer dependencies.

Surprisingly, sparse subsets of latents ($<$1% of total, e.g., $0.8\%$ for family, $0.9\%$ for function) recover up to 79% and 74% of original task performance, respectively. These compressed circuits define the core computational subroutines for each task, supporting direct interpretability and hypothesis generation.

Figure 2: ProtoMech as a replacement model for ESM2. It identifies task-relevant, interpretable latents and reconstructs critical circuit elements underlying predictive accuracy.

Further, in low-performance regimes (e.g., families where ESM2's F1$<$0.5), ProtoMech often denoises and surpasses the original model—implying that circuit tracing and sparse discrimination can enhance robustness by filtering context-specific noise.

Figure 3: Performance breakdown for protein family classification on challenging tasks (F1<0.5), highlighting ProtoMech’s denoising effect.

Mechanistic Interpretability: Biological Motifs and Mutation Sensitivity

ProtoMech's circuits exhibit strong alignment with established biochemical knowledge. Case studies demonstrate interpretable cross-layer pathways:

• Protein kinase domain circuits decompose into motifs detecting catalytic arginines (early), HRD catalytic loop (mid-layers), and ATP-binding/Glycine-rich loops (deep layers).
• NADP+ binding domain circuits capture residues participating in Rossmann fold formation and cofactor pockets, with motif reiteration in terminal layers—mirroring observations from mechanistic LLM studies.

  Figure 4: Examples of mechanistically coherent protein family circuits, mapping latents to known catalytic or binding motifs.

Mutation sensitivity analysis on GB1 highlights ProtoMech's ability to trace causal impact: high-fitness variants (e.g., A24W) activate stabilizing/residue-packing latents, whereas deleterious mutations (W43Q) inactivate both binding and hydrophobic interaction latents.

Figure 5: Mutation sensitivity in GB1, showing latent circuit shifts that reflect experimentally validated fitness changes.

Steering and Protein Design

ProtoMech supports targeted sequence “steering” by clamping key latent activations in the CLT surrogate, introducing directed mutational bias toward high-fitness regions. On seven DMS assays, ProtoMech-generated sequences routinely outperform those generated by random mutation or global concept-based steering (Contrastive Activation Addition, CAA).

Figure 6: Performance comparison of replacement models (ProtoMech, PLT, CAA) on DMS steering; circuit-based steering consistently yields higher functional scores.

ProtoMech-designed sequences generate the top-fitness variant in 71% of cases, and systematically populate upper deciles of fitness, indicating precise navigation of the protein fitness landscape and implying practical utility for mechanistically informed protein engineering.

Visual Analytics and Interactive Circuit Exploration

The ProtoMech visualizer operationalizes these findings, providing:

• Sequence-positioned latent activation grids (Figure 7)
• Layer-wise latent rankings (Figure 8)
• Sequence activation/alignments and influences for individual features (Figures 11–13)
• Construction and export of annotated circuit diagrams (Figure 9)
• Emphasis on cross-layer edge weights, facilitating topological circuit analysis.

  Figure 7: Main interface of the ProtoMech visualizer displaying cross-layer latent activation structure.

  

  Figure 9: Interactive circuit canvas, supporting custom annotation, analysis, and export of discovered circuits.

Theoretical Implications and Future Perspectives

ProtoMech substantiates that transformer-based pLMs operate via deeply interleaved, low-dimensional computational pathways, not merely layer-local feature detectors or shallow composites. Circuit tracing via CLTs provides a mathematically principled tool for extracting causal structure—transcending the static factorization of SAEs (Figure 10), and approaching the attribution graph paradigm now prominent in LLM interpretability.

Beyond interpretability, circuit-based surrogates can serve as regularizers, boosting model robustness by denoising internal representations, and as generators for high-throughput protein screening workflows by prioritizing sequence variants with validated mechanistic plausibility.

Figure 10: Conceptual distinction between SAEs (layer-local, reconstructive) and transcoders (cross-layer, computational).

Nevertheless, scaling CLT-based approaches faces quadratic parameter scaling in $L$ (number of layers), and automated annotation of novel circuit motifs remains an open challenge—potential avenues for continued methodological development.

Conclusion

ProtoMech introduces a robust framework for mechanistically tracing, interpreting, and steering computational circuits in protein LLMs. By constructing high-fidelity cross-layer transcoders that compress complex model behavior into interpretable, sparse circuits, ProtoMech enables both upstream scientific discovery and downstream engineering applications in computational biology. The empirical results establish new baselines for circuit fidelity, steerability, and biological alignment, catalyzing future research at the intersection of deep learning mechanistic interpretability and systems molecular biology.