Circuit Tracing in Autoregressive Protein Language Models

2026-06-14

AI Summary

This work addresses the lack of interpretability in autoregressive protein language models, which obscures their cross-layer computational dynamics. To this end, we propose ProGenMech, a framework that introduces cross-layer transcoders (CLTs) into autoregressive protein generation for the first time, enabling mechanistic interpretation of ProGen3’s generative and functional prediction processes through reconstruction of sparse latent variables. By integrating sparse mixture-of-experts (MoE) architectures with zero-shot circuit discovery, ProGenMech effectively identifies interpretable latent circuits associated with protein function. Experimental results demonstrate that ProGenMech outperforms baseline methods in causal generation and zero-shot fitness prediction tasks, accurately reproduces the original model’s output distribution, and successfully localizes biologically meaningful functional motifs and evolutionarily conserved sequence regions.

Abstract

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

Problem

• autoregressive generation
• protein language models
• mechanistic interpretability
• generative computation
• circuit tracing

Innovation

• mechanistic interpretability
• cross-layer transcoders
• autoregressive protein language models
• sparse latent circuits
• Mixture-of-Experts