Circuit Tracing in Autoregressive Protein Language Models

Abstract: Protein LLMs (pLMs) can generate novel protein sequences with properties beyond those observed in nature, yet the mechanisms underlying protein generation remain poorly understood. Existing mechanistic interpretability methods based on sparse autoencoders and transcoders primarily focus on protein representation learning models and do not capture the computation required for autoregressive generation. Here, we introduce ProGenMech, a mechanistic interpretability framework for generative protein LLMs that extends cross-layer transcoders (CLTs) to ProGen3, a sparse Mixture-of-Experts model trained for both causal generation and span infilling. Unlike per-layer approaches, CLTs reconstruct each layer using sparse latent variables from all preceding layers, enabling faithful recovery of inter-layer generative computation. We further develop a zero-shot circuit discovery framework to identify sparse latent circuits responsible for protein generation and fitness prediction. In causal generation and zero-shot fitness estimation tasks, ProGenMech outperforms local transcoder baselines in recovering ProGen3's probability distribution and functional scoring behavior, while matching the original model's generative distribution in span infilling tasks. Moreover, the recovered circuits reveal biologically meaningful motifs and functional regions associated with conserved sequence patterns and protein fitness landscapes, establishing a foundation for interpretable and steerable protein generation.

• The paper introduces ProGenMech, a CLT framework that exposes sparse, task-specific circuits in autoregressive protein language models.
• It demonstrates improved fidelity in sequence generation and fitness prediction using zero-shot circuit discovery with significant circuit compression.
• The study leverages attribution methods to map biological motifs, facilitating controlled protein design and enhanced model interpretability.

Circuit Tracing in Autoregressive Protein LLMs

Introduction

The study presents a mechanistic interpretability framework, ProGenMech, targeting autoregressive protein LLMs (pLMs)—specifically, ProGen3, a sparse Mixture-of-Experts (MoE) generative transformer optimized for both causal (CLM) and span-infill (GLM) protein sequence modeling. The novelty is rooted in extending cross-layer transcoders (CLTs) to generative architectures and establishing zero-shot circuit discovery protocols. This enables exposure and manipulation of interpretable computational subgraphs governing both sequence generation and protein fitness prediction.

Figure 1: ProGenMech as a sparse, circuit-level replacement for ProGen3, highlighting attribution-based discovery of functional motifs and improved log-likelihood alignment with the teacher model.

Cross-layer Transcoder Framework for Sparse MoE Generative Models

Modern sparse autoencoders (SAEs) elucidate interpretable latent features from pLMs but are limited in their capacity to expose computation distributed over multiple layers. Conventional per-layer transcoders (PLT) similarly suffer from limited expressivity, reconstructing layers independently and disregarding cross-depth computational accumulation. CLTs overcome these deficiencies by reconstructing each transformer/MLP layer as a function of sparse latents spread across all preceding layers, thus capturing the true trajectory of information flow and task-relevant computation throughout the network.

ProGenMech adapts this paradigm to the MoE regime, compressing all expert-path computation within each MoE layer into a common cross-layer latent space and altering training regimes for ProGen3’s hybrid causal/infill objectives. This involves a carefully balanced training protocol mirroring the teacher’s generative/infill sequence diversity and expert routing patterns.

Figure 2: ProGenMech CLT architecture, integrating across-layer latent bottlenecks with MoE blocks and simultaneous CLM/GLM training modalities.

Zero-shot Circuit Discovery

To bridge the gap between representation learning and mechanistic circuit discovery, ProGenMech deploys an attribution-based greedy search to identify sparse subcircuits of high-importance latent variables. These subcircuits are discovered in a zero-shot regime: the objective is to approximate the KL divergence (for generation) or Spearman correlation (for fitness ranking) between the reconstructed and original model outputs, only including as many latents as necessary.

The circuit search attributes credit to latents via product-activation-times-gradient scoring, incrementally including batches of high-contribution latents and halting at a specified fidelity threshold, typically when the sparse circuit recapitulates ≥70–95% of the teacher model’s predictive performance.

Numerical Results: Fidelity of Generation and Fitness Recovery

ProGenMech demonstrates strong fidelity versus baseline PLT approaches. With all latents, ProGenMech recovers approximately 60% of the teacher ProGen3’s generation likelihood (NLL), compared to 50–55% for PLT. Critically, highly compressed circuits—frequently using <2% of the latent space—retain 80–95% of teacher fitness prediction performance, emphasizing the redundancy and compressibility of functional computation in ProGen3. In GLM (span-infill) tasks, performance parity is observed between all models, attributed to the limited functional capacity of the relatively small ProGen3-112M backbone in GLM, as evidenced by elevated ground-truth GLM NLLs.

Figure 3: NLL and fitness recovery metrics for CLT (ProGenMech), PLT, and other replacement models in both generation and zero-shot settings.

Moreover, the mean number of latents per circuit exhibits a U-shaped pattern, with heavy usage in the initial/final layers and maximal sparsity in the intermediate layers, a trend echoing previously reported monosemanticity patterns in deep transformer interpretability.

Figure 4: Circuit compression—average number of latents per task, highlighting circuit sparsity and interpretability.

Biological Circuit Attribution and Visualization

A highlight of the study is the demonstration that CLT-located latents, composed into circuits, correspond to robust, biologically meaningful features:

• For protein kinases, distinct circuits are found for the HRD motif (catalytic loop) and DFG motif (ATP site), with upstream latents encoding single-residue detectors and downstream latents assembling broader motifs (Figure 5).
• In the GRB2 fitness landscape, wild-type and mutant analyses reveal interpretable attribution shifts in latents corresponding to binding interfaces and stability sites, aligning with experimentally established protein biophysics (Figure 6).

  Figure 5: CLM generation circuits in kinase, showing sequential motif detection from single amino acid level to full functional regions.

  

  Figure 6: Attribution comparison for GRB2 wild-type, high-fitness, and low-fitness variants, demonstrating interpretable redistribution of functional motifs.

Visual Analytics Infrastructure

The authors deploy a comprehensive interactive visualization toolkit supporting latent, sequence, and circuit graph inspection. This interface enables:

• Grid-wise mapping of latent activations by sequence position and model layer (Figure 7)
• Exploration of layer-wise latent rankings, per-latent sequence activation profiles, and motif alignments with reference proteomes (Figures 6–8)
• Interactive annotation and subgraph analysis of latent-latent and latent-output dependencies on a canvas interface (Figures 9–10)
• Direct attribution tracing from upstream latents through virtual-weight edges to final generation logits (Figure 8)

  Figure 7: Grid interface mapping layer- and position-specific latent activations, sequence axes, and connection overlays.

Implications and Directions for Future Mechanistic Interpretability

The success of ProGenMech CLTs in compressing and disentangling the computation of generative protein transformers extends mechanistic interpretability from masked representation backbones to the more complex generative setting. By exposing biologically meaningful, task-specific feature circuits, this approach:

• Opens the possibility for interpretability-driven protein sequence design via circuit steering or intervention
• Enables analysis of which motifs are leveraged during generation versus scoring, offering insight into gaps or misalignment between generation and downstream function prediction
• Sets groundwork for scaling: larger architectures demand improved parameter-efficient cross-layer mechanisms, possibly integrating expert-specific crosscoders for MoE models

Current limitations include incomplete steering of generative outputs toward desired fitness phenotypes (likely due to underpowered base models and fitness–generation misalignments) and a bottleneck in automating biological annotation of discovered subcircuits.

Conclusion

ProGenMech advances mechanistic interpretability for generative protein transformers by extending CLT-based sparse factorization, enabling recovery and analysis of task-relevant computational circuits for both generation and fitness prediction. Numerical benchmarks demonstrate improved fidelity and dramatic circuit compression, while qualitative analyses map learned computations to known biological motifs. This work substantively extends the interpretability toolkit for protein LLMs and signals key avenues for controlled and explainable protein sequence design.

---

Reference: "Circuit Tracing in Autoregressive Protein LLMs" (2606.16044)