File size: 9,994 Bytes
f4cc4c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 | """Hidden-state sparse autoencoders for ESM++ (ESMC) models.
FastPLMs implements the published Biohub hidden-state SAE contract directly, so attaching an SAE
needs only PyTorch, Transformers, and the checkpoint itself. Biohub still owns the SAE weights:
this module reads their published repository layout, one shared ``config.json`` plus one
``layer_{index}.safetensors`` shard per backbone layer, and never redistributes those tensors.
A layer produced here satisfies the same attachment contract as an official Biohub
``ESMCSAEModel.layers`` entry, so ``PreTrainedESMplusplusModel.add_sae_models`` accepts either.
"""
from __future__ import annotations
import json
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from torch import Tensor
SAE_CONFIG_FILE = "config.json"
SAE_REQUIRED_CONFIG_FIELDS = ("d_model", "codebook_dim", "k")
# Trained per-feature statistics. Shards that never trained them omit them, and the ones default
# makes the (features / max) * idf normalization an identity.
SAE_OPTIONAL_STATE_NAMES = ("idf", "max")
STANDARDIZATION_EPSILON = 1e-5
@dataclass(frozen=True, slots=True)
class ESMplusplusSAEParams:
"""Shape contract of one hidden-state SAE.
``layer`` indexes the ESM++ hidden state the SAE reads, where ``0`` is the embedding output and
``num_hidden_layers`` is the final normalized state.
"""
d_model: int
codebook_dim: int
k: int
layer: int
@dataclass(slots=True)
class ESMplusplusSAEOutput:
"""Sparse features, and the optional reconstruction error, for one batch of residues."""
feature_magnitudes: Tensor
reconstruction_loss: Tensor | None = None
def _standardize_residue_states(residue_states: Tensor) -> Tensor:
"""Center and scale each residue vector the way the SAEs were trained."""
# residue_states: (n, d)
centered = residue_states - residue_states.mean(dim=-1, keepdim=True) # (n, d)
return centered / (centered.std(dim=-1, keepdim=True) + STANDARDIZATION_EPSILON) # (n, d)
class ESMplusplusSAELayer(nn.Module):
"""Top-k sparse autoencoder over one ESM++ hidden state.
Encoding standardizes each residue vector, projects it into a wide codebook, and keeps only the
``k`` largest activations. The decoder is used only to measure reconstruction error, which is
why it is opt-in: interpretation and gradient-based design read ``feature_magnitudes`` alone.
"""
idf: Tensor
max: Tensor
def __init__(self, params: ESMplusplusSAEParams) -> None:
super().__init__()
self.params = params
self.W_enc = nn.Parameter(torch.empty(params.d_model, params.codebook_dim)) # (d, c)
self.W_dec = nn.Parameter(torch.empty(params.codebook_dim, params.d_model)) # (c, d)
self.b_dec = nn.Parameter(torch.zeros(params.d_model)) # (d,)
self.register_buffer("idf", torch.ones(params.codebook_dim)) # (c,)
self.register_buffer("max", torch.ones(params.codebook_dim)) # (c,)
@property
def layer(self) -> int:
"""ESM++ hidden-state index this SAE was trained against."""
return self.params.layer
def forward(
self,
residue_states: Tensor,
*,
compute_reconstruction_loss: bool = False,
) -> ESMplusplusSAEOutput:
# residue_states: (n, d) for n residues; c is the codebook width, k the retained features
standardized = _standardize_residue_states(residue_states) # (n, d)
preactivations = F.relu((standardized - self.b_dec) @ self.W_enc) # (n, c)
retained = torch.topk(preactivations, self.params.k, dim=-1) # values, indices: (n, k)
feature_magnitudes = torch.zeros_like(preactivations).scatter(
-1, retained.indices, retained.values
) # (n, c)
if not compute_reconstruction_loss:
return ESMplusplusSAEOutput(feature_magnitudes=feature_magnitudes)
reconstructed = feature_magnitudes @ self.W_dec + self.b_dec # (n, d)
return ESMplusplusSAEOutput(
feature_magnitudes=feature_magnitudes,
reconstruction_loss=(reconstructed - standardized).pow(2).mean(dim=-1), # (n,)
)
def get_sae_output(self, layer_states: Tensor, token_mask: Tensor) -> ESMplusplusSAEOutput:
"""Encode the unpadded residues of one ESM++ hidden state.
This name and signature are the attachment contract that
``PreTrainedESMplusplusModel.add_sae_models`` validates and calls.
"""
# layer_states: (b, l, d); token_mask: (b, l)
residue_states = layer_states[token_mask] # (n, d) for n valid tokens
encoded: ESMplusplusSAEOutput = self(residue_states)
return encoded
def _repository_file(
repository: str | os.PathLike[str],
filename: str,
*,
revision: str | None,
cache_dir: str | os.PathLike[str] | None,
token: str | bool | None,
local_files_only: bool,
) -> Path:
"""Resolve one repository file, from a local directory or the Hub cache."""
# A directory counts as local only when it holds the shared config, so a stale directory named
# like a Hub identifier cannot shadow the download.
local_directory = Path(repository)
if (local_directory / SAE_CONFIG_FILE).is_file():
path = local_directory / filename
if not path.is_file():
raise FileNotFoundError(f"SAE repository {local_directory} has no {filename}.")
return path
return Path(
hf_hub_download(
repo_id=str(repository),
filename=filename,
revision=revision,
cache_dir=None if cache_dir is None else str(cache_dir),
token=token,
local_files_only=local_files_only,
)
)
def _load_sae_layer(
shard_path: Path,
*,
params: ESMplusplusSAEParams,
device: torch.device,
dtype: torch.dtype | None,
) -> ESMplusplusSAELayer:
state = load_file(str(shard_path), device=str(device))
encoder = state.get("W_enc")
if encoder is None:
raise ValueError(f"{shard_path} is not an ESMC SAE shard; it has no 'W_enc' entry.")
# Build on the meta device so the shard tensors are the only materialized copy of a codebook
# that reaches roughly one gigabyte for the widest published SAEs.
with torch.device("meta"):
sae_layer = ESMplusplusSAELayer(params)
sae_layer.to(dtype=encoder.dtype if dtype is None else dtype)
sae_layer.to_empty(device=device)
# to_empty leaves the statistics buffers uninitialized, so restore the identity defaults that
# shards without trained statistics rely on.
sae_layer.idf.fill_(1.0)
sae_layer.max.fill_(1.0)
incompatible = sae_layer.load_state_dict(state, strict=False)
missing = tuple(
name for name in incompatible.missing_keys if name not in SAE_OPTIONAL_STATE_NAMES
)
unexpected = tuple(incompatible.unexpected_keys)
if missing or unexpected:
raise ValueError(
f"{shard_path} does not match the ESMC SAE state contract; "
f"missing {list(missing)}, unexpected {list(unexpected)}."
)
return sae_layer
def load_esmc_sae_layers(
repository: str | os.PathLike[str],
layers: Sequence[int],
*,
revision: str | None = None,
cache_dir: str | os.PathLike[str] | None = None,
token: str | bool | None = None,
local_files_only: bool = False,
device: torch.device | str = "cpu",
dtype: torch.dtype | None = None,
) -> dict[int, ESMplusplusSAELayer]:
"""Load the requested hidden-state SAE layers from a Hub repository or a local directory.
Only the shared config and the requested shards are read, so a repository that publishes every
backbone layer costs one shard per requested layer. ``dtype`` defaults to the dtype stored in
the shard; pass the ESM++ model dtype when the SAE has to consume its hidden states directly.
"""
requested = tuple(dict.fromkeys(int(layer) for layer in layers))
if not requested:
raise ValueError("Loading SAE layers requires at least one backbone layer index.")
config_path = _repository_file(
repository,
SAE_CONFIG_FILE,
revision=revision,
cache_dir=cache_dir,
token=token,
local_files_only=local_files_only,
)
config = json.loads(config_path.read_text(encoding="utf-8"))
missing_fields = [name for name in SAE_REQUIRED_CONFIG_FIELDS if name not in config]
if missing_fields:
raise ValueError(f"{config_path} is not an ESMC SAE config; it omits {missing_fields}.")
available = tuple(int(index) for index in config.get("available_layers", ()))
target_device = torch.device(device)
sae_layers: dict[int, ESMplusplusSAELayer] = {}
for layer in requested:
if available and layer not in available:
raise ValueError(
f"SAE repository {repository} does not publish layer {layer}; "
f"available layers are {list(available)}."
)
shard_path = _repository_file(
repository,
f"layer_{layer}.safetensors",
revision=revision,
cache_dir=cache_dir,
token=token,
local_files_only=local_files_only,
)
sae_layers[layer] = _load_sae_layer(
shard_path,
params=ESMplusplusSAEParams(
d_model=int(config["d_model"]),
codebook_dim=int(config["codebook_dim"]),
k=int(config["k"]),
layer=layer,
),
device=target_device,
dtype=dtype,
)
return sae_layers
__all__ = [
"ESMplusplusSAELayer",
"ESMplusplusSAEOutput",
"ESMplusplusSAEParams",
"load_esmc_sae_layers",
]
|