Protein Circuit Tracing via Cross-layer Transcoders

AI summary

Overview

Research area: Mechanistic interpretability of protein language models (pLMs), specifically circuit discovery and steering in transformer-based protein models, with applications to protein design.

Technical level: Advanced. The paper assumes familiarity with transformer internals (residual streams, MLP blocks, attention), sparse autoencoders, transcoders, and protein fitness benchmarks.

Scope: This paper introduces ProtoMech, a framework that uses cross-layer transcoders (CLTs) to build a replacement model of the protein language model ESM2, discovers compact interpretable circuits inside it, and uses those circuits for protein fitness steering and biological motif analysis.

What This Paper Is About

Protein language models like ESM2 predict protein structure and function well, but the internal computations they use to make those predictions are largely opaque. Existing interpretability methods such as sparse autoencoders decompose a single layer's activations into interpretable features, but they do not model the layer-to-layer transformations that constitute the model's actual computation. ProtoMech addresses this by learning sparse latent representations jointly across layers so that each MLP layer's output is reconstructed from the latents of all preceding layers, producing a replacement model whose circuits can be traced, pruned, and steered.

Key Contributions

1. ProtoMech framework. The authors develop a framework for discovering and analyzing computational circuits within and across transformer layers in protein language models, built on cross-layer transcoders that predict each layer's output from the sparse latents of all preceding layers. Applied to ESM2, it achieves what the paper describes as state-of-the-art recovery of the original model's performance: 89% on protein family classification and 82% on function prediction.

2. Compact circuit discovery. The authors identify highly compressed circuits retaining 79% (family) and 74% (function) of model performance while using less than 1% of ProtoMech's latent space. The results section separately reports 79% recovery for family circuits and 76% recovery for function circuits. Steering along these circuits generates diverse functional sequences, including the top-fitness variants in 71% of cases.

3. Biological validation. Through case studies of kinase activity, NADP+ binding, and the GB1 fitness landscape, the authors show that ProtoMech recovers computational pathways aligning with known structural and functional motifs, including binding, signaling, and stability.

4. Open-source release. ProtoMech is released as open-source software with accompanying visualization tools (code at github.com/amirgroup-codes/ProtoMech, visualizer at protmech.github.io).

Main Findings

• Replacement model fidelity: Using the full set of latents, ProtoMech achieves an average F1 of 0.82 ± 0.19 on protein family classification, recovering approximately 89% of ESM2's performance (0.92 ± 0.10), versus a PLT baseline F1 of 0.50 ± 0.34. On function prediction it achieves an average Spearman correlation of 0.41 ± 0.19, recovering about 82% of the original model's performance (0.50 ± 0.22), versus 0.38 ± 0.18 for the PLT baseline.

• Circuit compression: ProtoMech circuits reach an average F1 of 0.73 ± 0.19 on family classification (79% recovery) using only 150 ± 204 latents on average, roughly 0.8% of the total latent space; the PLT circuit reaches 0.49 ± 0.33. On function prediction, circuits reach an average Spearman of 0.38 ± 0.18 (76% recovery) using 167 ± 221 latents, about 0.9% of the latent space; the PLT circuit reaches 0.35 ± 0.19.

• Robustness across mutation depth: On the GFP_AEQVI_Sarkisyan DMS assay, which contains variants up to 10 mutations away, ProtoMech circuits maintain predictive power across all mutation depths, recovering up to 74% of ESM2's predictive power at 5+ mutations.

• Steering outperforms alternatives: On seven of the 12 DMS assays, steering via ProtoMech and PLT circuits consistently yields higher-fitness sequences than Contrastive Activation Addition (CAA) and random mutation selection. ProtoMech outperforms PLT in 71% of cases, including generating the single highest-fitness variant and the highest top 10% and top 20% fitness variants.

• Interpretable motifs in family circuits: A kinase domain circuit (InterPro ID IPR000719) shows layer 1 latent L1/1582 activating on arginine, feeding into L3/1918 which detects the conserved HRD catalytic motif; layer 5 splits into L5/1774 (ATP-binding site) and L5/826 (glycine-rich loop), and L6/886 activates on serine. An NADP+ binding circuit (InterPro ID IPR036291) shows L1/967 activating on phenylalanine, feeding L4/53 (Rossmann fold), then L5/1338 (NADP+ binding pockets), with L5/1109 detecting NADP+ and FAD-binding sites and L6/2198 reactivating phenylalanine.

• Mutation sensitivity in GB1: For the wildtype GB1 sequence, L1/249 activates on tryptophan at position 43 (W43), feeding L3/3027 (hydrophobic interactions, notably between leucine at position 5 and W43) and L6/1610 (tryptophan detection). The high-fitness variant A24W retains the wildtype circuit and adds L6/3029, associated with inward-facing and hydrophobic residues. The low-fitness variant W43Q deactivates the W detectors L1/249 and L6/1610 and weakens L6/3027.

• Denoising effect: Discovered circuits occasionally outperform full ESM2. On family classification tasks where ESM2 scores below 0.5 F1 (0.39 ± 0.04), ProtoMech circuits achieve 0.43 ± 0.27 on average; in fitness prediction, 37% of the circuits achieve higher performance than ESM2.

• Scaling to ESM2-35M: ProtoMech circuits recovered up to 85% of ESM2-35M's performance, compared to up to 79% for ESM2-8M. A "windowed" CLT variant reduced parameter count by 40% (207M to 125M) and sped up training by 1.75x, capturing 82% of ESM2-35M's performance versus 68% for the PLT baseline and 85% for the vanilla CLT.

Methodology in Plain English

The authors replace each MLP block inside ESM2 with a sparse bottleneck. An encoder maps a layer's residual stream activation into a large latent vector, keeping only the top-k largest activations (a TopK function). A decoder then reconstructs the MLP output, but crucially it does so using decoders that map latents from all preceding layers to the current layer, not just from the current layer. The transcoders are trained by minimizing mean squared error between the original and reconstructed MLP outputs across all layers, plus an auxiliary loss (inspired by prior work) that reactivates "dead" latents by decoding the top-k_aux latents.

With this replacement model in hand, the authors treat the task-specific circuit as the minimal set of latents whose joint activity reproduces the model's behavior. They first train a probe (logistic regression for family classification, a CNN for fitness prediction) on the original final-layer MLP output to establish a performance reference. Then they apply an iterative greedy search driven by gradient-based attribution: latents are ranked by their contribution to the probe's output on a held-out validation set and added incrementally until the circuit reaches at least 70% of the original ESM2 performance, or matches the full replacement model. Evaluation uses F1 for family classification and Spearman rank correlation for fitness prediction.

For steering, they clamp selected latents in the target circuit to a high activation value (the maximum activation magnitude observed for that node across the sequence multiplied by a scalar), compute the final-layer reconstruction, and decode to ESM2's logits, selecting mutations with the maximum probabilities. Generated variants are constrained to within five mutations of the wildtype, and fitness is independently assessed with a CNN evaluation model trained on 90% of the training data. For qualitative analysis, the visualizer builds sparse graphs where nodes are latents and edges are weighted by the product of a source node's activation and the gradient of the target's pre-activation with respect to the source node; activated amino acids are cross-referenced with the top-10 maximally activating sequences from Swiss-Prot and projected onto protein structures.

Training data: 5 million protein sequences of length up to 1022 amino acids randomly sampled from UniRef50. ESM2-8M has L = 6 layers and d_model = 320; ESM2-35M has L = 12 layers and d_model = 480. Family tasks use Swiss-Prot sequences clustered at 30% sequence identity with InterPro family labels; fitness tasks use 12 DMS assays from ProteinGym. On ESM2-8M, the CLT used k = 16 per layer (total sparsity 16 × 6 = 96 latents), d_latent = 3200 (a 10x expansion), k_aux = 32, and alpha = 1/32. Training a CLT on ESM2-8M required approximately 28M parameters, a 3.5x increase over the original model's size.

Why This Matters

Impact on research: The paper argues that SAEs only factorize representations and cannot serve as replacement models, while per-layer transcoders (PLTs) approximate each layer in isolation and neglect accumulation of computation across depth, leading to degraded representations and unreliable circuit recovery. ProtoMech's cross-layer formulation provides an empirical demonstration that modeling inter-layer dependencies improves both fidelity to the original model and the quality of extracted circuits, which is a methodological shift for mechanistic interpretability in biology.

Real-world applications:

• High-throughput protein screening: using ProtoMech as a "mechanistic filter" to prioritize variants that rely on biologically plausible motifs for wet-lab synthesis.
• Protein design and engineering: steering discovered circuits to generate high-fitness variants, as demonstrated on seven DMS assays.
• Model regularization: the observed denoising effect suggests sparse transcoding could improve robustness of pLM predictions on downstream tasks.
• Interpretable diagnostics: providing mechanistic rationales for why particular mutations raise or lower fitness, as illustrated by the GB1 A24W and W43Q analyses.

Industry relevance: Any organization using protein language models for therapeutic design, enzyme engineering, or directed evolution pipelines could use circuit-level tracing to screen candidate variants more efficiently and to audit whether a model's predictions rest on biologically meaningful features. The open-source release and visualizer reduce the barrier to adoption.

Future Directions

• Automated circuit annotation: The authors state that interpretation currently relies on manual analysis plus existing biological annotations, that labelling capacity is bounded by current biological knowledge, and that automated annotation pipelines are a critical next step.
• Extending to other pLM architectures: The paper notes that extending ProtoMech to autoregressive and diffusion-based pLMs is not trivial, since autoregressive models often need alignment or fine-tuning strategies that differ from the supervised probing used here, and CLTs have not been applied to diffusion models to the authors' knowledge.
• Scaling and efficiency: Because CLT decoder matrices scale as O(L²), training larger pLMs is costly; "windowed" CLTs are proposed as a tradeoff, and the paper reports that circuits on ESM2-35M recovered up to 85% of performance versus up to 79% on ESM2-8M, suggesting further scaling may be increasingly effective.
• Mitigating interpretation limits: Developing ways to label circuits beyond current biological knowledge and to scale qualitative analysis beyond human-dependent review.

Target Audience

Researchers and practitioners in mechanistic interpretability, protein language modeling, and computational protein design; machine learning scientists interested in replacement models and transcoder architectures; and bioinformatics or protein engineering groups that want interpretable, steerable tools for fitness prediction and variant screening. Readers should have working familiarity with transformer internals and protein fitness benchmarks such as DMS assays and ProteinGym.

Authors’ abstract

Protein language models (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.