mjkmain/language_balanced_sae

Language-balanced sparse autoencoders

TopK sparse autoencoders trained on raw transformer-block outputs from Qwen/Qwen3-1.7B-Base, using a balanced total of 700 million retained training tokens: 100 million each for English, Chinese, French, Spanish, Vietnamese, Indonesian, and Japanese. Access to model files requires manual approval.

Choose a checkpoint

Property	Values
Latent width	4096, 8192, 16384, 32768, 65536, 131072
TopK	50 (ReLU before TopK; actual nonzero count can be lower)
Transformer block index, zero-based	3, 8, 13, 17, 22, 26
Input dimension	2048
Saved training checkpoints	14 per width, approximately every 50M tokens
Default checkpoint	tokens_000700000000 (final)
Total SAEs	504, including 36 final SAEs

catalog.json lists every SAE, its exact training token count, weight path, configuration, and validation results. Filter is_final: true for the final models. All checkpoint weights together occupy approximately 355.2 GB; the complete collection including configurations and evaluation results is approximately 360.9 GB. The final 36 SAE weight files occupy approximately 25.4 GB.

Repository layout

README.md
catalog.json
manifest.json
w4096_k50/                       # one directory per latent width
  tokens_000700000000/           # actual retained training tokens
    config.json                 # training settings and provenance
    layers.3/
      cfg.json                  # native sparsify SAE configuration
      sae.safetensors           # native sparsify weights
    layers.8/ ...
    validation/                 # metrics and feature frequencies
  tokens_000050028059/ ...       # intermediate checkpoints
  eval_step_500/ ...             # evaluations between checkpoints
  timing.json
w8192_k50/ ...

Paths preserve the original training layout. Intermediate checkpoint token counts are the actual counts at optimizer-step boundaries, rather than rounded labels. training_state.pt (optimizer and training-resume state) is excluded. Raw training data and credentials are not part of this repository. manifest.json records the relative paths and byte sizes of all source files included in the upload.

Download one SAE

Request access on the model page and wait for approval, then authenticate with hf auth login. Use a read token for downloads. Install huggingface_hub, safetensors, and torch in your Python environment.

from pathlib import Path
from huggingface_hub import snapshot_download

repo_id = "mjkmain/language_balanced_sae"
sae_id = "w16384_k50/tokens_000700000000/layers.13"
root = snapshot_download(
    repo_id,
    allow_patterns=[f"{sae_id}/cfg.json", f"{sae_id}/sae.safetensors"],
    token=True,
    # revision="<commit SHA>",  # pin a revision for reproducible experiments
)
sae_dir = Path(root) / sae_id

If using the project's vendored sparsify implementation:

from sparsify import SparseCoder

sae = SparseCoder.load_from_disk(sae_dir, device="cpu").eval()

Alternatively, the following plain PyTorch example uses the native weights without requiring sparsify. x must contain raw outputs of the selected block, with shape [n_tokens, 2048].

import json
import torch
import torch.nn.functional as F
from safetensors.torch import load_file

cfg = json.loads((sae_dir / "cfg.json").read_text())
weights = load_file(str(sae_dir / "sae.safetensors"), device="cpu")

@torch.no_grad()
def encode_decode(x):
    x = x.to(weights["encoder.weight"])
    pre = F.linear(x - weights["b_dec"], weights["encoder.weight"],
                   weights["encoder.bias"]).relu()
    values, indices = pre.topk(cfg["k"], dim=-1, sorted=False)
    reconstruction = F.embedding_bag(
        indices, weights["W_dec"], per_sample_weights=values, mode="sum"
    ) + weights["b_dec"]
    return values, indices, reconstruction

The native encoder subtracts b_dec before the linear layer. Decoder weights have shape [width, 2048]. Do not apply input normalization or rescaling. These files use native sparsify semantics, not the bias-folded Qwen-Scope export format.

Training provenance and evaluation

The base-model revision is ea980cb0a6c2ae4b936e82123acc929f1cec04c1. English training data comes from HuggingFaceFW/fineweb; the other six languages come from HuggingFaceFW/fineweb-2. Full dataset revisions, masks, training-order hashes, hyperparameters, and the upstream sparsify commit are recorded in each checkpoint's config.json. The implementation includes local patches listed in that file; an upstream commit alone does not reproduce the complete training code.

Training excludes sequence position 0, EOS tokens, shared norm outliers, and tokens beyond each language's quota. Inputs are not normalized. Language balance refers to retained token counts; it does not imply balanced topics or concepts.

Checkpoint validation/ directories contain per-language metrics and feature frequencies. Frequency sidecars record language order and interpretation. eval_step_* directories contain additional evaluation snapshots. FLORES is designated evaluation-only in the training provenance.