VFUSE: Virulent Feature Understanding With Sparse AutoEncoders

Abstract Generative models have shown remarkable progress in a variety of domains such as protein design, but such power enables the opaque generation of hazardous proteins. In this work, we introduce VFUSE (Virulent Feature Understanding with Sparse auto Encoders), a mechanistic interpretability approach that trains SAEs on diffusion-transformer activations to audit protein models for hazard-aware features. We apply VFUSE to Rose TTAFold3 and RFDiffusion3, popular open-weight models for protein folding and synthesis. We find that for certain blocks, linear probes detect hazardous designs significantly better when fit in the SAE latent space over the original model's representations: improving interpretability without sacrificing model performance. Furthermore, we identify monosemantic features from the SAE that fire only on hazardous designs at up to AUROC 0.84. To our knowledge this is the first SAE trained on an all-atom diffusion model and the first feature-level virulence audit of a protein design model, paving the way towards safe and interpretable protein design.BlogThere has been a ton of mechanistic interpretability research done on LLMs, from SAELens to Neuronpedia to Golden Gate Claude. Even CNNs and ViTs seem to have a bunch of interesting work. An area that seems relatively underinvestigated is protein model interpretability. There are some early papers here such as InterProt and FoldSAE, but so many unanswered questions and possibilities. We wanted to answer the question: Can SAEs (Sparse Autoencoders) trained on RFDiffusion3 and RoseTTAFold3 be used to to classify hazardous vs non-hazardous proteins in an interpretable way?We trained Matryoshka Batch TopK Sparse Autoencoders (SAEs) on diffusion transformer activations in RFDiffusion3 (RFD3, a generative protein model) and RoseTTAFold3 (RF3, a protein structure prediction model like AlphaFold), sampling 1475 length-matched benign/hazardous pairs from UniProt/SafeProtein + ToxinPred3.To simulate generation