evolve-away/Boltz1-SAEs-L2-rec1

📄 Model Card for Boltz-1 SAEs: Pairformer Trunk (Recycle 1)

This repository contains the trained TopK Sparse Autoencoders (SAEs) for the Pairformer Trunk evaluated at Recycle Iteration 1 as featured in the preprint “Where a folding model keeps biology: probing and sparse-autoencoder analysis of the Boltz-1 trunk and diffusion module”.

These dictionaries map the dense representation spaces of the main trunk into an interpretable, sparse latent basis.

Model Details

• Base Architecture: Boltz-1 (Open-source AlphaFold3-class structure predictor)
• SAE Type: TopK Sparse Autoencoder ($k=256$, total latents $N=2048$)
• Regularization: $L_2$ weight regularization ($3 \times 10^{-3}$) on training-set demeaned activations.

Architectural Scope: Recycle 1

This repository hosts dictionaries trained explicitly within the Pairformer Trunk during its final execution loop:

• Layers Available: Layers 0 to 47
• Recycling Configuration: Recycle 1. This represents the main text's core focus, acting as the final interface representation fed directly into the downstream diffusion decoder.

Training Dataset & Preprocessing

The SAEs were trained via unsupervised dictionary learning using a structural biology activation dataset of 84,074 unlabelled proteins (~21.96M residues total).

• Demeaning Step: Prior to encoding, the training-set mean activation was subtracted. This counteracts Boltz-1's extreme activation-energy concentration (where ~98% of raw squared magnitude sits in the top 10 dimensions), reducing the top-10 energy share to ~29% and enabling the dictionary to learn fine, distributed biology.

Key Scientific Insights (Trunk Stack)

If you use these models, please contextualize them with the architectural insights established in our paper:

• Semantic Compiler: The Pairformer trunk linearly encodes both macro-geometry (DSSP secondary structure F1 0.79–0.90) and highly localized sequence chemistry (signal peptides, disordered regions, and disulfide bonds with a targeted mid-trunk hotspot at fractional depth 0.2–0.35).
• SAE Feature Concentration: While supervised linear probes dominate overall decodability, these SAEs offer clean, monosemantic single-feature alignment for amino-acid identities (e.g., Cysteine latent F1 $\approx 0.99$) where raw individual neurons fail.