This repository contains the trained TopK Sparse Autoencoders (SAEs) for the Diffusion Coordinate Module evaluated 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 activation spaces of Boltz-1's generative coordinate decoder 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$)
This specific repository hosts dictionaries trained explicitly within the Diffusion Coordinate Module:
Layers Available: Layers 0 to 22
Sampling Trajectory Steps: Dictionaries are provided across specific time-steps of the generative denoising trajectory:
step_0 (Highest noise initialization)
step_1, step_10, step_50, step_100
step_199 (Final denoised 3D structural coordinates)
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 (Diffusion Stack)
If you use these models, please contextualize them with the architectural insights established in our paper:
Spatial Coordinate Engine: The diffusion module keeps macro-geometric features entirely intact but rapidly attenuates sequence chemistry along both the layer and sampling-step axes (e.g., Signal peptide probe F1 drops $0.76 \to 0.35$; Cysteine identity drops $1.0 \to 0.66$).
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.
