ProtoMech: Protein Circuit Tracing Framework

• ProtoMech Framework is a method for mechanistic interpretability in pLMs, using sparse cross-layer transcoders to reconstruct transformer MLP pathways.
• It introduces sparse latent spaces that enable high-fidelity circuit recovery and efficient protein family classification with extreme model compression.
• The framework facilitates targeted protein design by steering top-attribution latent circuits, aligning model outputs with known biological motifs.

ProtoMech Framework

ProtoMech is a framework developed for mechanistic interpretability and circuit tracing in protein LLMs (pLMs), specifically addressing the challenge of identifying and manipulating the internal computational circuits that underpin the predictive success of models such as ESM2. ProtoMech introduces sparse cross-layer transcoders capable of reconstructing the full multi-layer perceptron (MLP) pathway of a transformer-based pLM, achieving high-fidelity recovery of function and family classification performance with extreme compression and distinct interpretability advantages. The framework enables effective steering of model output for protein design, establishes correspondence between circuit structure and known biological motifs, and provides a blueprint for principled mechanistic dissection of deep learning models in protein informatics (Tsui et al., 12 Feb 2026).

1. Motivation and Problem Definition

Modern pLMs exhibit strong performance across protein structure and function prediction tasks, but their predictions remain mechanistically opaque. Existing interpretability methods such as sparse autoencoders and per-layer transcoders operate at individual layers, failing to capture the cross-layer computational dependencies critical for understanding the model’s end-to-end inference. ProtoMech defines a "computational circuit" as the minimal set of pathways traversing layers, positions, and channels whose joint activity drives downstream predictions (e.g., family classification, function regression). Capturing such circuits requires modeling the explicit dependencies across all transformer layers, as each MLP block in a pLM builds recursively on the outputs of all preceding layers.

2. ProtoMech Architecture

ProtoMech replaces the original stack of MLP blocks in a transformer pLM with a sparse cross-layer transcoder (CLT), introducing a bottlenecked latent space per layer and a decoder that reconstructs each layer's output from all previous layers’ latent codes. This architecture is parameterized as follows:

• For layer $\ell$, $x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$ denotes the residual-stream input, and $y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$ is the corresponding MLP output.
• ProtoMech encodes $x^{(\ell)}$ as a sparse latent $a^{(\ell)} \in \mathbb{R}^{d_\text{latent}}$:

$$a^{(\ell)} = \text{TopK}(W_\text{enc}^{(\ell)} \cdot (x^{(\ell)} - b_\text{pre}^{(\ell)}) + b_\text{enc}^{(\ell)})$$

Here, TopK zeroes all but $k \ll d_\text{latent}$ largest-magnitude entries.

• The cross-layer decoder reconstructs $\hat{y}^{(\ell)}$ as:

$$\hat{y}^{(\ell)} = \sum_{\ell'=1}^{\ell} W_\text{dec}^{(\ell' \rightarrow \ell)} \cdot a^{(\ell')} + b_\text{pre}^{(\ell)}$$

• Global encoder-decoder notation:

$$z = E(\{ x^{(\ell)} \}_{\ell=1}^{L}; W_E), \quad \hat{x}^{(\ell')} = D(z; W_D^{(\ell')})$$

• For ESM2-8M: $x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$0, $x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$1, $x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$2, $x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$3.

This cross-layer construction enables each layer to incorporate information from all latents activated in its own and all previous layers, supporting the tracing of distributed computational circuits.

3. Training Objective and Regularization

ProtoMech's loss function comprises:

• Layerwise mean-squared error reconstruction:

$x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$4

• Sparsity penalty (optional $x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$5 over latents, though hard sparsity is imposed by TopK):

$x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$6

• Auxiliary loss to avoid dead latents and diversify circuit usage:

$x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$7

where $x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$8, and $x^{(\ell)} \in \mathbb{R}^{d_\text{model}}$9 is the top $y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$0 decoder reconstruction.

The final training loss is

$y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$1

with $y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$2 (e.g., family classification or DMS regression probe), and $y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$3 tuned for fidelity and utility.

4. Training Protocol and Empirical Performance

ProtoMech was trained on ESM2-8M using 5M UniRef50 sequences (≤1022 residues). Training details:

• TopK sparsity: $y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$4 per layer, $y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$5, $y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$6.
• Adam optimizer, learning rate $y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$7, batch size 16, weight decay $y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$8, 250k steps.
• Baseline: per-layer transcoder (PLT) reconstructs each $y^{(\ell)} = \text{MLP}^{(\ell)}(x^{(\ell)})$9 from $x^{(\ell)}$0 alone (no cross-layer reconstruction).
• Evaluation: replaced MLPs with either CLT or PLT in a sequential replacement setup; ground-truth attention fixed.

Performance metrics (held-out test set):

Task	ProtoMech CLT	Original ESM2	PLT Baseline
Family F1 ($x^{(\ell)}$1)	0.82 ± 0.19	0.92 ± 0.10	0.50 ± 0.34
Function Spearman ($x^{(\ell)}$2)	0.41 ± 0.19	0.50 ± 0.22	0.38 ± 0.18

Circuit compression: greedy accumulation of top-attribution latents until ≥70% of ESM2’s metric is recovered yields <1% of the latent space (150–167 latents on average), with ProtoMech achieving 74–79% recovery. PLT circuits with the same budget yield only 35–49% recovery.

5. Mechanistic Circuit Interpretability

Visualization and attribution of ProtoMech’s sparse circuits reveal mechanistically relevant correspondences to known protein motifs, exemplified as follows:

• Protein kinase domain (IPR000719):
  • L1/1582: Arg sensor (arginine building blocks).
  • L3/1918: HRD catalytic loop (His-Arg-Asp).
  • L5/1774: ATP-binding pocket.
  • L5/826: Glycine-rich loop (G-loop).
  • L6/886: Serine site (phosphorylation).
• NADP⁺ binding domain (IPR036291):
  • L1/967: Phenylalanine in binding pocket.
  • L4/53: Rossmann fold ($x^{(\ell)}$3 motif).
  • L5/1338/1109: NADP⁺, FAD binding.
  • L6/2198: Phe motif recapture.
• GB1 IgG-binding domain circuits tracked mutational effects: circuits adapt or ablate specific latents corresponding to fitness-altering mutations.

Circuit structure is visualized as directed graphs of latent units (nodes) and inter-latent influences ("virtual weights" $x^{(\ell)}$4), explicating activation and inhibition patterns across the computation.

6. Steering and Protein Design Applications

ProtoMech permits targeted intervention for protein design by clamping activation of top-attribution latents identified from functional circuits. The procedure:

• Sample $x^{(\ell)}$5 top-attribution latents ($x^{(\ell)}$6), scaling their activation at high-response residues in a wildtype sequence.
• Decode through layer $x^{(\ell)}$7's cross-layer decoder to generate perturbed $x^{(\ell)}$8, add any residual correction, and decode the logits for variant selection.
• Up to five mutations (positions with maximum logit change) are selected for experimental validation.

Empirical benchmarks across seven deep-mutational scanning (DMS) assays indicate that ProtoMech circuit steering yields higher mean and maximum fitness gains over PLT, CAA, and random baselines, achieving superior outcomes in 71% of settings. For example, in GFP, ProtoMech achieves mean fitness +4.17 (random +2.74) and max fitness +4.63 (+3.54), and consistently finds the highest-fitness variants.

7. Implications, Limitations, and Future Directions

ProtoMech bridges sparse feature factorization and per-layer surrogate approaches by furnishing a full multi-layer replacement that closely emulates transformer MLP computation, while revealing tractable and interpretable computational circuits. The sparse bottleneck effect can filter noise and, in low-performance regimes (e.g., families where ESM2 F1 < 0.5), ProtoMech sometimes surpasses the original pLM.

Applications encompass mechanistic filters for high-throughput protein screening and the design of biologically plausible mutational strategies anchored in explicit circuit diagrams. Noted limitations include $x^{(\ell)}$9 decoder parameter scaling (28M parameters for CLT vs. 8M in ESM2-8M) and the reliance on manual circuit annotation. Prospective research directions include (a) scaling to larger pLMs through efficient cross-layer compression, (b) automation of circuit annotation via structure-mapping and motif databases, (c) extension to attention-head circuits, multi-task probes, and unsupervised circuit discovery, and (d) integration of ProtoMech as a regularizer in pLM fine-tuning or in-the-loop wet-lab design workflows (Tsui et al., 12 Feb 2026).