FoldSAE: Steerable Protein Structure Design

• FoldSAE is a mechanistic framework that employs sparse autoencoders to extract interpretable secondary structure features from RFdiffusion activations.
• The approach uses block-level ablations and logistic probing to identify helix and strand detectors, enabling tunable interventions during protein design.
• Steering experiments demonstrate that controlled adjustments in activations can modulate helical and strand content while maintaining biological plausibility.

FoldSAE is a mechanistic interpretability framework for RFdiffusion, a generative model of protein structures, which employs sparse autoencoders (SAEs) to discover interpretable, secondary structure-specific features within the model's activations and leverages these features for direct, tunable control over secondary structure formation during protein generation (Zarzecki et al., 27 Nov 2025). The approach establishes a connection between internal neural network representations and secondary structure content, providing protein designers with precise steering capabilities over helix and strand content while retaining the biological plausibility of generated proteins.

1. Background and Rationale

RFdiffusion repurposes the RoseTTAFold network for denoising in a diffusion-probabilistic model over protein backbones. A protein of length $L$ is represented as $L$ rigid frames $x = [x_1,…,x_L]$, with each frame $x_l = (z_l, r_l)$ comprising $z_l \in \mathbb{R}^3$ (C$\alpha$ coordinates) and $r_l \in SO(3)$ (N–C$\alpha$–C orientation). Generation starts from fully noised coordinates and iteratively applies one of 36 stacked "blocks" (main and extra, with global-attention), culminating in a backbone suitable for sequence design (e.g., via ProteinMPNN).

Despite RFdiffusion’s empirical capabilities, its internal mechanisms governing secondary structure formation are opaque. Mechanistic interpretability, as demonstrated in prior SAE work, decomposes dense activations into sparse, often interpretable features. FoldSAE’s innovation is the application of SAEs to RFdiffusion block activations to unsupervisedly identify helix/strand “detectors” and expose handles for controlling generation output.

2. SAE Block Selection and Model Architecture

Identification of the block most informative for secondary structure is performed by ablation: for each block index $m$, compare original RFdiffusion $M_{\text{orig}}$ to a variant $L$0 where block $L$1 is replaced by block $L$2. The index $L$3 that induces the largest change in overall helical content, as measured by a summary statistic $L$4, is selected:

$L$5

Empirically, $L$6 "main_04", with all subsequent SAE modeling using this block's activations.

The activations (size $L$7) are flattened and split into $L$8 per-residue patches $L$9. The SAE comprises:

• Encoder: single fully-connected layer $x = [x_1,…,x_L]$0, bias $x = [x_1,…,x_L]$1, ReLU.
• TopK sparsification: retains only $x = [x_1,…,x_L]$2 largest positive entries per code $x = [x_1,…,x_L]$3.
• Decoder: single linear layer $x = [x_1,…,x_L]$4.

Formally,

$x = [x_1,…,x_L]$5

With expansion factor 16 and $x = [x_1,…,x_L]$6, $x = [x_1,…,x_L]$7 latent units; $x = [x_1,…,x_L]$8 active per residue.

Training minimizes $x = [x_1,…,x_L]$9, with TopK enforcing sparsity (no explicit $x_l = (z_l, r_l)$0 penalty). The trained model achieves $x_l = (z_l, r_l)$1 explained variance, minimal proportion of “dead” units (<$x_l = (z_l, r_l)$2), and low high-frequency activation rates (<$x_l = (z_l, r_l)$3).

3. Feature Discovery and Steering Protocol

To identify secondary structure features, 10,000 unsteered backbones are generated; per-residue encodings $x_l = (z_l, r_l)$4 are labeled as Helix/Strand/Coil with STRIDE. One-vs-rest logistic probes on $x_l = (z_l, r_l)$5 achieve approximately 84% balanced accuracy and >93% ROC AUC. Units with high positive coefficients for helix vs rest and high negative coefficients for strand vs rest (or vice versa) are classified as “detectors.” Approximately 50–100 of $x_l = (z_l, r_l)$6 units show substantial probe coefficients ($x_l = (z_l, r_l)$7), indicating monosemantic, antagonistic features for helix and strand.

The steering protocol introduces a continuous hyperparameter $x_l = (z_l, r_l)$8:

• $x_l = (z_l, r_l)$9: no intervention,
• $z_l \in \mathbb{R}^3$0: steer toward target (helix or strand),
• $z_l \in \mathbb{R}^3$1: steer away.

At each diffusion step, chosen block activations $z_l \in \mathbb{R}^3$2 are encoded as $z_l \in \mathbb{R}^3$3; for each detector $z_l \in \mathbb{R}^3$4, assign sign $z_l \in \mathbb{R}^3$5 indicating target association. Each active $z_l \in \mathbb{R}^3$6 is modulated:

$z_l \in \mathbb{R}^3$7

The decoder reconstructs modified activations $z_l \in \mathbb{R}^3$8. To counteract reconstruction error, an error offset $z_l \in \mathbb{R}^3$9 (computed at $\alpha$0) is added, yielding $\alpha$1; denoising proceeds with these intervened activations.

Steering for $\alpha$2-helices sets $\alpha$3 for helix detectors and $\alpha$4 for strand detectors, with $\alpha$5; strand control reverses signs or uses $\alpha$6. Changes propagate, altering local $\alpha$7 torsions associated with secondary structures.

4. Interpretability and Internal Structure Analysis

Block-level ablations confirm “main_04” as carrying non-redundant information for helix/strand formation, with its removal collapsing helical content nearly to zero. Analysis across the $\alpha$8 latent units reveals that only a sparse subset are meaningfully active in association with secondary structure.

Activation patterns, as $\alpha$9 increases, show 30–50% elevation in helix-detector firing on residues originally predicted as coil or strand, reliably predicting coil$r_l \in SO(3)$0helix transitions. Conversely, negative $r_l \in SO(3)$1 on helix detectors produces coil$r_l \in SO(3)$2strand transitions, demonstrating interpretable, antagonistic control across two distinct motifs at both the neuron and block level.

5. Evaluation: Steering Efficiency and Structural Fidelity

Quantitative evaluation involves sweeping $r_l \in SO(3)$3 from –5 to +5:

• Helix-targeted steering increases helix content from approximately 50% to 60% at $r_l \in SO(3)$4; further gains are limited by intrinsic saturation.
• Strand-targeted steering raises strand content from ~30% to ~55% at $r_l \in SO(3)$5, with coil content remaining largely stable.

Case studies show progressive, localized transitions in single-protein backbones under $r_l \in SO(3)$6 variation, visually confirming residue-level control. After steering, sequence design with ProteinMPNN followed by sequence–embedding metrics (FBD, MMD) demonstrates that generated proteins remain within $r_l \in SO(3)$7 of the unsteered baseline, indicating no significant loss in biological plausibility.

6. Implications and Limitations for Protein Design

FoldSAE enables direct modulation of secondary structure content via human-understandable controls, reducing reliance on post hoc filtering of random samples and allowing for mechanistic, designer-driven intervention. The framework provides evidence that monosemantic, sparse features corresponding to structural motifs are accessible and actionable within large-scale generative models of protein structure.

Current work focuses on secondary-structure steering. Proposed extensions include targeting other blocks or network tracks (e.g., 1D/2D features for solvent exposure or pocket geometry), designing regionally restricted interventions (e.g., residues 1–20 only), and employing additional regularizers (e.g., $r_l \in SO(3)$8, KL, total-variation) to mediate sparsity–reconstruction tradeoffs.

A recognized limitation is the evaluation bottleneck; sequence plausibility currently requires a round-trip through ProteinMPNN, imposing latency. The use of direct geometric or energetic structural metrics (e.g., Ramachandran outliers, hydrogen-bond criteria) presents a potential avenue for accelerating validation. FoldSAE establishes a paradigm for interpretable, steerable protein backbone design via mechanistically grounded neural feature manipulation within diffusion models (Zarzecki et al., 27 Nov 2025).