| --- |
| tags: |
| - mechanistic-interpretability |
| - biology |
| - sparse-autoencoders |
| - protein-structure |
| --- |
| |
| # 📄 Model Card for Boltz-1 SAEs: Diffusion Module |
|
|
| 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$) |
| - **Regularization:** $L_2$ weight regularization ($3 \times 10^{-3}$) on training-set demeaned activations. |
| |
| ## Architectural Scope: Diffusion Module Only |
| 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. |
