living-models/Botanic1-M-sae-k64-codebook8192

Botanic1-M SAEs — k=64, codebook 8192

BatchTopK sparse autoencoders over the residual stream of living-models/Botanic1-M, one per probed block. Each decomposes a base's 1024-d hidden state into a sparse combination of 8192 learned features.

Layer	Depth	Held-out EV	Threshold	rms_scale	Dead features
13	0.25	0.830	1.0119261741638184	1.6462319514670176	0.66%
26	0.50	0.780	0.966841995716095	0.5425183788529638	0
35	0.67	0.749	0.8141039609909058	0.2214300936095105	0
47	0.90	0.727	0.881266176700592	0.13001471952870194	0

Shared across layers: d_in 1024 → d_sae 8192 (expansion 8), BatchTopK k=64, site hidden_states[layer] — the output of that BiMamba2 block. Full trainer metrics (L0 distribution, activation-frequency deciles, held-out MSE) are in metrics.json.

Usage

import torch
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("living-models/Botanic1-M", trust_remote_code=True).eval()
tokenizer = AutoTokenizer.from_pretrained("living-models/Botanic1-M", trust_remote_code=True)
sae = AutoModel.from_pretrained(
    "living-models/Botanic1-M-sae-k64-codebook8192", trust_remote_code=True
).eval()
print(sae.available_layers)   # [13, 26, 35, 47]

# PHYB (AT2G18790) in Arabidopsis thaliana: 500 bp around the ATG start codon.
# 5'UTR is bases 0-249; the CDS starts at base 250.
sequence = (
    "TTTTTTTTTGTTATCTCTCTCTATCTGAGAGGCACACATTTTGCTTCGTCTTCTTCAATTTATTTTATTGGTTTCTC"
    "CACTTATCTCCGATCTCAATTCTCCCCATTTTCTTCTTCCTCAAGTTCAAAATTCTTGAGAATTTAGCTCTACCAGA"
    "ATTCGTCTCCGATAACTAGTGGATGATGATTCACCCTAAATCCTTCCTTGTCTCAAGGTAATTCTGAGAAATTTCTC"
    "AAATTCAAAATCAAACGGCATGGTTTCCGGAGTCGGGGGTAGTGGCGGTGGCCGTGGCGGTGGCCGTGGCGGAGAA"
    "GAAGAACCGTCGTCAAGTCACACTCCTAATAACCGAAGAGGAGGAGAACAAGCTCAATCGTCGGGAACGAAATCTC"
    "TCAGACCAAGAAGCAACACTGAATCAATGAGCAAAGCAATTCAACAGTACACCGTCGACGCAAGACTCCACGCCGT"
    "TTTCGAACAATCCGGCGAATCAGGGAAATCATTCGACTACT"
)

z = sae.encode_sequences([sequence], model, tokenizer, layer=35)  # (1, n_bases, 8192)
print(z.shape, "active per base:", (z[0] > 0).sum(-1).float().mean().item())

encode_sequences tokenizes, runs the backbone, takes hidden_states[layer], drops the <cls> column and encodes — so z[b, i] is the feature vector of base sequence[i].

Which features fire, and where

peak, base = z[0].max(dim=0)                      # each feature's strongest base
top = peak.topk(10)
for value, feature in zip(top.values, top.indices):
    print(f"feature {feature.item():5d}  peak {value.item():7.3f}  at base {base[feature].item()}")
# feature  1589  peak  22.790  at base 250   <- the A of the ATG start codon

# the strongest features at the A of the ATG start codon
atg = z[0, 250].topk(5)
print(list(zip(atg.indices.tolist(), atg.values.tolist())))

Comparing depths

for layer in sae.available_layers:
    z = sae.encode_sequences([sequence], model, tokenizer, layer=layer)
    peak, base = z[0].max(dim=0)
    strongest = int(peak.argmax())
    print(f"layer {layer:>2}  L0 {(z[0] > 0).sum(-1).float().mean():6.2f}  "
          f"top feature {strongest:5d} @ base {base[strongest].item()}")

Encoding activations you already have

with torch.no_grad():
    out = model(input_ids=tokenizer([sequence], return_tensors="pt")["input_ids"],
                output_hidden_states=True, return_logits=False)
    hidden = out.hidden_states[35][:, sae.config.cls_offset:]   # (1, n_bases, 1024)

    z = sae.encode(hidden, layer=35)        # applies rms_scale + the pinned threshold
    recon = sae.decode(z, layer=35) / sae.config.layer_config["35"]["rms_scale"]

print("fraction of variance kept:", (1 - (recon - hidden).var() / hidden.var()).item())
# 0.8791

Reading a feature's direction

dictionary = sae.dictionary(35)
direction = dictionary.W_dec[2810]     # (1024,) — feature 2810 in residual space
cosine = torch.nn.functional.normalize(dictionary.W_dec, dim=-1) @ \
         torch.nn.functional.normalize(direction, dim=-1)
print("nearest features:", cosine.topk(6).indices.tolist()[1:])

Notes

Four things change the numbers if you bypass encode_sequences:

• rms_scale is per layer, and the spread is wide — 1.65 at layer 13 down to 0.13 at layer 47. Pairing one layer's activations with another layer's scale is a silent order-of-magnitude error, so encode takes layer and reads the scale itself rather than accepting one.
• <cls> shifts every coordinate. The Botanic1 tokenizer prepends <cls> (tokenizer("ACGT")["input_ids"] == [3, 4, 6, 7, 5]), so base i is token i + 1. Slice it off before encoding.
• hidden_states[i] is the output of block i. Botanic1's list has exactly 52 entries and no embedding entry, so there is no off-by-one to correct here — hidden_states[35] is the right tensor for the layer-35 dictionary.
• Batch equal lengths. Botanic1's SSM scans do not mask pad tokens, so a padded batch changes the activations of the real bases in shorter rows.

These dictionaries were fit on activations captured in bfloat16, while the released Botanic1-M runs in float32. Feature activations from this repository therefore agree with the published feature statistics to about 0.4% median relative error and 98-99% support overlap, not exactly, at every layer.

encode applies the single threshold the trainer converged on rather than a per-batch top-k, so one position's activations never depend on which other positions travelled with it.

Citation

Technical report: BOTANIC-1: a series of long-context plant genomic foundation models in the agentic era (bioRxiv, 2026).

@article{Barozet2026.09.04.749355,
    author = {Barozet, Am{\'e}lie and Cabeli, Vincent and Ogier du Terrail, Jean
              and Rukhovich, Alexey and Janssoone, Thomas and Klajer, Gary
              and Sheikhitarghi, Zeinab and Andrews, Gregory and Veran, Cyril
              and Strouk, L{\'e}onard},
    title = {BOTANIC-1: a series of long-context plant genomic foundation
             models in the agentic era},
    journal = {bioRxiv},
    year = {2026},
    elocation-id = {2026.09.04.749355},
    doi = {10.64898/2026.09.04.749355},
    publisher = {Cold Spring Harbor Laboratory},
    URL = {https://www.biorxiv.org/content/early/2026/09/09/2026.09.04.749355},
    eprint = {https://www.biorxiv.org/content/early/2026/09/09/2026.09.04.749355.full.pdf},
}

Research use only

See LICENSE. Not for production, clinical or diagnostic use.