🧬 biology

VFUSE: Virulent Feature Understanding with Sparse autoEncoders

This paper introduces VFUSE, a mechanistic interpretability framework that employs sparse autoencoders on diffusion-transformer activations to identify and audit hazardous features in protein design models like RoseTTAFold3 and RFDiffusion3, achieving high-accuracy detection of virulent designs while maintaining model performance.

Original authors: Michael Yu, Matthew L. Olson

Published 2026-06-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Michael Yu, Matthew L. Olson

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you have a super-smart robot chef that can invent new recipes for proteins (the building blocks of life). This robot is amazing; it can create enzymes that clean oil spills or binders that catch viruses. But, like any powerful tool, it has a dark side: if you ask it to build something dangerous, like a poison, it might do so without you even realizing it until it's too late.

Currently, security guards for these robots only check the final "shopping list" (the sequence of ingredients) to see if it looks familiar to known poisons. If the list looks new, the guard lets it pass, even if the final dish is toxic. They can't explain why it's dangerous, and they can't stop the robot while it's cooking.

Enter VFUSE: The "X-Ray Goggles" for Protein Robots

This paper introduces a new tool called VFUSE. Think of VFUSE as a pair of X-ray goggles that let you see exactly what the robot is thinking while it's designing the protein, not just what it produced at the end.

Here is how it works, broken down into simple steps:

1. The "Translator" (The Sparse Autoencoder)

The robot chef (specifically models called RoseTTAFold3 and RFDiffusion3) thinks in a complex, messy language that humans can't understand. It's like a brain where thousands of neurons fire at once, mixing up ideas about shape, texture, and danger all together.

VFUSE uses a special translator called a Sparse Autoencoder (SAE). Imagine taking that messy brain and sorting every single thought into its own separate, labeled drawer.

Before VFUSE: The robot's brain is a giant, tangled ball of yarn.
After VFUSE: VFUSE untangles the yarn and puts each strand into its own clear box. Now, we can see exactly which "thought" corresponds to "danger" and which corresponds to "safety."

2. The "Spotlight" (Finding the Danger)

The researchers trained this translator on the robot's internal thoughts. They found that at a specific stage of the robot's thinking process (called "Block 12"), the translator works best.

When they looked through the sorted drawers, they found specific "danger switches."

The Analogy: Imagine a room full of people talking. You can't hear who is saying what. VFUSE puts a spotlight on specific people.
The Result: They found that certain "switches" in the robot's brain light up only when the robot is designing a poison. For example, one switch lit up when the robot was designing a snake toxin, and another lit up for a virus-evasion protein.
Precision: These switches don't just say "Poison!" for the whole protein. They point to the exact few ingredients (amino acids) that make it dangerous, like highlighting the specific spice that makes a dish toxic.

3. The "Memory Test" (Did it Cheat?)

A common trick in AI is "memorization." If you train a student on a list of 100 poisons, they might just memorize those 100 names rather than learning what makes something poisonous.

The researchers tested VFUSE to see if it was just memorizing. They gave the robot new, slightly different versions of known poisons (like a cousin of a snake toxin).

The Result: VFUSE still spotted the danger. This proves the tool actually learned the concept of virulence, not just a list of names. In fact, VFUSE was better at spotting these "cousin" poisons than looking at the robot's raw, unsorted thoughts.

4. Why This Matters (According to the Paper)

The paper claims this is the first time anyone has done this kind of "thought-reading" on these specific, high-tech protein-design robots.

Better Security: It can detect dangerous designs better than current methods that only look at the final ingredient list.
Explainability: It doesn't just say "Stop!" It shows you where in the design the danger is hiding (e.g., "It's in this specific spiral shape on the surface").
No Performance Loss: Using this tool to check for danger didn't slow down the robot or make it worse at making good proteins.

In Summary:
VFUSE is like installing a transparent dashboard in a self-driving car. Instead of waiting for the car to crash to see if it was going the wrong way, VFUSE lets you see the car's internal map in real-time, highlighting exactly which turn it's planning to take that leads to a cliff, allowing you to intervene before the accident happens.

Technical Summary: VFUSE – Virulent Feature Understanding with Sparse autoEncoders

Problem Statement

Generative models for protein design, particularly all-atom diffusion models like RoseTTAFold3 (RF3) and RFDiffusion3 (RFD3), possess the dual-use capability to design hazardous proteins (e.g., neurotoxins, immune evasion proteins) by scaffolding new structures around dangerous templates. Current biosecurity screening methods rely on sequence homology or sequence-level classifiers. These approaches suffer from two critical limitations: they operate post-hoc (flagging sequences after generation) and lack mechanistic interpretability, meaning they cannot explain which specific structural features triggered a hazard flag or intervene during the sampling process. Furthermore, while Sparse Autoencoders (SAEs) have been successfully applied to protein language models (PLMs) and vision models, they have not yet been utilized to probe the internal representations of all-atom diffusion models.

Methodology

The authors introduce VFUSE, a mechanistic interpretability framework that trains Sparse Autoencoders on the activations of diffusion-transformer models to audit for hazard-associated features.

Model & Activation Collection:
- The study targets RFDiffusion3 (an 18-block all-atom diffusion model) and RoseTTAFold3 (its folding counterpart).
- Activations ( $d=768$ ) are captured via PyTorch hooks at specific residual stream blocks: blocks 6, 8, and 12 for RFD3, and blocks 12 and 16 for RF3.
- For RFD3, inputs are generated via partial diffusion (at $\partial t = 5$ , approx. 5Å noise) on PDB structures of hazardous templates to simulate the model's design process. RF3 inputs are folded directly.
Dataset Construction:
- Virulent Sequences: Sourced from SafeProtein (275 pairs) and ToxinPred3 (1,200 sequences for RF3 only).
- Benign Sequences: Length-matched (100–300 aa) sequences from UniProt, filtered to exclude toxin/virulence keywords.
- Homology Control: To prevent probes from memorizing fold-family identity rather than learning virulence signals, sequences are clustered at 30% identity using mmseqs2. Train/test splits ensure no near-homologs appear in both sets.
SAE Architecture & Training:
- Architecture: A Matryoshka BatchTopK SAE is employed. This combines:
  - Matryoshka: Nested prefix sizes allowing reconstruction at multiple granularities.
  - BatchTopK: Enforces sparsity by keeping only the top $K$ pre-activations across a batch, replacing L1 penalties.
- Hyperparameters: Dictionary size $m=12,288$ (16x expansion), $K=80$ active latents, trained for 20,000 steps.
- Performance: The trained SAEs explain 96.9% of held-out activation variance with an average $L_0$ sparsity of 79.9.
Evaluation Protocol:
- Probing: Mean-pooled SAE latent vectors ( $z$ ) and raw activations ( $x$ ) are used to train L2-regularized logistic regression probes.
- Metrics: AUROC is evaluated under two regimes: random stratified splits and homology-clustered splits.
- Feature Scoring: Univariate AUROC is computed for each feature against class labels. Significance is determined via Mann–Whitney U tests with Benjamini–Hochberg FDR correction ( $q < 0.05$ ).

Key Results

SAE Latents Outperform Raw Activations:
- At RFD3 Block 12, SAE-encoded features significantly outperform raw activations in the homology-clustered setting, achieving an AUROC of 0.817 compared to 0.763 for raw features (a +0.054 improvement).
- This improvement is specific to Block 12; earlier blocks (6, 8) show neutral or negative gains from SAEs.
- For RF3 (ToxinPred3 dataset), Block 12 SAE features also show a +0.041 improvement over raw features in clustered splits.
Reduction in Homology Memorization:
- The gap between random-split and homology-clustered AUROC (a proxy for family memorization) is smallest for RFD3 Block 12 SAE features (0.13 gap) compared to Block 6 (0.22 gap). This suggests Block 12 SAEs capture class-discriminative, family-generalizable structural signals rather than memorizing fold families.
Identification of Monosemantic Hazard Features:
- The study identifies individual SAE features that fire selectively on hazardous designs.
- Top Feature: Feature 639 in RFD3 Block 12 achieves an AUROC of 0.819 ( $q \approx 10^{-13}$ ) for detecting viper neurotoxins.
- Localization: These features exhibit per-residue localization. For Feature 639, the top-3 activating residues (99, 104, 109) cluster within a single alpha-helix of the Ammodytoxin A structure. Other features similarly localize to specific regions of Monkeypox immune evasion proteins, ribosome-inactivating proteins, and mosquito salivary toxins.
Competitive Performance:
- The best VFUSE probe (RFD3 Block 12 SAE, random split, AUROC 0.877) is competitive with state-of-the-art sequence-only classifiers (e.g., DTVF at 0.92), despite using a simple logistic probe and a smaller dataset. Crucially, VFUSE provides per-residue rationales, which sequence-only models lack.

Significance and Claims

The paper claims to be the first SAE trained on an all-atom diffusion model and the first feature-level virulence audit of a protein design model.

Mechanistic Insight: The work demonstrates that dictionary learning at specific depths (Block 12) untangles polysemantic structural concepts into monosemantic directions that are more amenable to low-capacity hazard detection than raw model representations.
Safety & Interpretability: VFUSE establishes a pathway for runtime monitoring and interpretable auditing of generative protein models. By identifying specific residues and structural motifs associated with virulence, the method moves beyond "black box" screening to explainable safety checks.
Future Direction: While the authors note limitations (small dataset, lack of steering success in preliminary tests), the work paves the way for safe and interpretable protein design by enabling the detection of hazardous features at the feature level rather than just the sequence level.