Advanced Usage

TopKSAETrainer is the easy path, but it is a thin wrapper. When you need a custom training loop, a different loss, a non-linear encoder, or direct control over the sparsification, drop down to the building blocks. This page covers the lower-level objects exported at the top level of compresso.

The raw model: `TopKSAE`

TopKSAE is a plain nn.Module. Its forward returns a (reconstruction, codes, stats) triple, where codes already has exactly k non-zeros per row and stats is a dict of monitoring metrics:

import torch
from compresso import TopKSAE

model = TopKSAE(input_dim=128, hidden_dim=512, k=32, tied=False)
x = torch.randn(256, 128)

reconstruction, codes, stats = model(x)
# stats keys: reconstruction_mse, cosine_similarity,
#             active_count, activation_freq, dead_features

Writing your own training loop is then completely standard PyTorch:

opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(50):
    perm = torch.randperm(x.size(0))
    for i in range(0, x.size(0), 128):
        batch = x[perm[i : i + 128]]
        _recon, _codes, stats = model(batch)
        loss = stats["reconstruction_mse"]
        opt.zero_grad()
        loss.backward()
        opt.step()

The `stats` dictionary

The metrics returned each forward pass are useful both as losses and as health checks:

reconstruction_mse: Mean squared error between input and reconstruction.
cosine_similarity: Mean per-row cosine similarity; the trainer optimizes a blend of (1 - cosine_similarity) and MSE (see alpha_loss).
active_count: Mean number of active features per row (equals k for a standard top-k SAE).
activation_freq: Per-feature firing rate over the batch, shape (hidden_dim,).
dead_features: Count of features that never fired in the batch. A large value means part of your dictionary is wasted — lower k, lower hidden_dim, or train longer.

Encoder, decoder, and tied weights

By default the encoder and decoder are single nn.Linear layers, but you can supply any modules — for example a deeper, non-linear encoder — as long as the shapes line up:

import torch.nn as nn
from compresso import TopKSAE

encoder = nn.Sequential(
    nn.Linear(784, 256), nn.GELU(), nn.Linear(256, 512),
)
model = TopKSAE(input_dim=784, hidden_dim=512, k=16, encoder=encoder)

Set tied=True to make the decoder reuse the encoder weight (transposed), which halves the parameter count and is common for SAEs. Use model.get_decoder_weight() to fetch the effective decoder matrix in either case (this is what the A First Example: Seeing What an SAE Learns plots as dictionary atoms).

Controlling sparsification

The bottleneck is a reusable layer, TopKSparsify, backed by the functional topk_ste(). You can drop either into any model:

from compresso import TopKSparsify, topk_ste

sparsify = TopKSparsify(k=8, score_mode="abs", ste_alpha=0.01)
z = sparsify(torch.randn(4, 64))     # exactly 8 non-zeros per row

z2 = topk_ste(torch.randn(4, 64), k=8, score_mode="abs", ste_alpha=0.01)

Two knobs matter:

score_mode selects which entries survive the top-k:
- "abs" keeps the largest-magnitude values (signed features; the default).
- "raw" keeps the largest signed values.
- "relu" keeps the largest positive values and discards negatives.
ste_alpha is the straight-through estimator leak. The forward pass is a hard top-k (non-differentiable), so the backward pass routes a fraction ste_alpha of the gradient to the non-selected entries and full gradient to the selected ones. ste_alpha=0 is a pure hard mask; a small value such as 0.01 keeps unused features learning and reduces dead features.

These are surfaced on the config as sparsify_score_mode / sparsify_ste_alpha and srp_score_mode (the latter is used by transform when packing into an SRPTensor).

Post-sparsification hooks

A post_sparsify module runs on the codes after the top-k. The built-in L1Normalize and L2Normalize rescale each code to unit L1/L2 norm, which is handy when codes feed a downstream similarity or retrieval step:

from compresso import TopKSAEConfig, TopKSAETrainer, L1Normalize

cfg = TopKSAEConfig(hidden_dim=4096, k=128, post_sparsify=L1Normalize())
trainer = TopKSAETrainer(cfg)

Full config reference

Every trainer hyperparameter lives on TopKSAEConfig:

Field	Default	Meaning
`hidden_dim`	`4096`	Number of dictionary features `H`.
`k`	`128`	Active features kept per row.
`decoder_bias`	`False`	Add a bias to the default decoder.
`pre_act`	`None`	Module applied before sparsification.
`post_sparsify`	`None`	Module applied to codes after top-k.
`encoder` / `decoder`	`None`	Custom modules (else linear layers).
`sparsify_score_mode`	`"abs"`	Top-k scoring: `abs` / `raw` / `relu`.
`sparsify_ste_alpha`	`0.01`	Straight-through leak for non-selected entries.
`alpha_loss`	`0.01`	Cosine/MSE mixture weight in the training loss.
`l1_penalty`	`0.0`	Extra L1 penalty on code activations.
`batch_size`	`128`	Rows per batch.
`shuffle`	`True`	Shuffle rows between epochs.
`seed`	`42`	Seed for shuffling and init.
`epochs`	`10`	Training epochs.
`lr` / `weight_decay`	`1e-3` / 0	AdamW parameters.
`decay`	`False`	Cosine LR decay to zero over training.
`compile`	`False`	`torch.compile` the model when available.
`device`	`"cpu"`	Training/transform device.
`show_progress`	`True`	tqdm progress bar when tqdm is installed.
`srp_score_mode`	`"abs"`	Score mode for `SRPTensor.from_dense` in transform.

Sparse parameters and pruning

Beyond representation learning, Compresso ships sparse parameter types for compressing model weights:

MaskedParam — a weight with a learned/scheduled binary mask for magnitude pruning.
SRPParam — a parameter backed by the same fixed-k sparse layout as SRPTensor.
SparsityController — a global dispatcher that advances and rewinds MaskedParam masks during training, and exponential_decay(), a helper for sparsity schedules.

Note

The pruning stack (and the broader compresso.layers package of sparse Linear/Embedding/attention layers) is experimental and not part of the stable first-release surface. The representation-learning API on this page and in Input and Output is the supported path; expect the parameter/pruning APIs to change.