Core API

Sparsification

compresso.topk_ste(x, k, dim=-1, score_mode='abs', ste_alpha=0.0)[source]

Hard top-k forward with alpha-scaled STE on non-selected entries.

Forward keeps exactly k signed values per slice (others are zero). Backward uses gradient scale 1.0 on selected top-k entries and ste_alpha on all non-selected entries.

Return type:

Tensor

Parameters:

x (Tensor)
k (int)
dim (int)
score_mode (str)
ste_alpha (float)

class compresso.TopKSparsify(k, dim=-1, score_mode='abs', ste_alpha=0.0)[source]

Activation layer that applies hard top-k sparsification with STE.

Parameters:

k (int) – Number of entries to keep along dim.
dim (int) – Dimension along which to select (default -1).
score_mode (str) – Scoring mode for top-k ranking (abs, raw, relu).
ste_alpha (float) – Backward gradient scale for non-selected positions.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Apply top-k sparsification to an input tensor.

The forward pass keeps the k highest-scoring entries along dim and sets all other entries to zero. Ranking is controlled by score_mode and gradients for non-selected entries are scaled by ste_alpha through the straight-through estimator.

Parameters:: x (Tensor) – Input tensor to sparsify.
Returns:: Tensor with the same shape as x and at most k non-zero entries along dim.
Return type:: Tensor

set_k(k)[source]

Change k at runtime (e.g. for scheduling).

Return type:: None
Parameters:: k (int)

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:: str

Sparse Representations

class compresso.SRPTensor(*, cols, vals, shape, prefix_shape=None, validate=True)[source]

Minimal SRP tensor container.

Stores a row-packed sparse matrix with fixed k nonzeros per row. The representation uses cols and vals tensors of shape (rows, k) with logical dense shape (rows, cols_total).

Optionally carries prefix_shape so callers can restore (*prefix, cols_total).

Parameters:

cols (Tensor)
vals (Tensor)
shape (Tuple[int, int])
prefix_shape (Tuple[int, ...] | None)
validate (bool)

select_rows(row_indices)[source]

Alias for gradient-preserving row indexing.

Return type:: SRPTensor
Parameters:: row_indices (slice | Tensor | Sequence[int])

to_dense()[source]

Densify to (rows, cols_total) or (*prefix, cols_total).

Return type:: Tensor

to_coo()[source]

Convert to a coalesced 2D PyTorch sparse COO tensor.

The returned tensor has logical shape (rows, cols_total). Duplicate SRP columns in a row are summed by COO coalescing, matching to_dense() scatter-add semantics.

Return type:: Tensor

to_csr()[source]

Convert to a 2D PyTorch sparse CSR tensor.

Return type:: Tensor

to_csc()[source]

Convert to a 2D PyTorch sparse CSC tensor.

Return type:: Tensor

to_bsr(blocksize)[source]

Convert to a 2D PyTorch sparse BSR tensor.

Parameters:: blocksize (tuple[int, int]) – Sparse block size (row_block, col_block). Both dimensions must divide the logical dense shape.
Return type:: Tensor

to_bsc(blocksize)[source]

Convert to a 2D PyTorch sparse BSC tensor.

Parameters:: blocksize (tuple[int, int]) – Sparse block size (row_block, col_block). Both dimensions must divide the logical dense shape.
Return type:: Tensor

to_scipy_coo()[source]

Convert to a SciPy coo_matrix on CPU.

This conversion detaches tensors and therefore does not preserve autograd history.

to_scipy_csr()[source]: Convert to a SciPy csr_matrix on CPU.

to_scipy_csc()[source]: Convert to a SciPy csc_matrix on CPU.

to_numpy_dict()[source]

Return a structural NumPy representation.

The returned dictionary contains cols, vals, shape, and prefix_shape. It is intentionally not dense; use srp.to_dense().cpu().numpy() when a dense NumPy array is desired.

Return type:: dict[str, Any]

numpy()[source]

Alias for to_numpy_dict().

SRPTensor is structurally represented by both cols and vals, so this method returns structural arrays rather than a dense matrix.

Return type:: dict[str, Any]

to_dict()[source]

Serialize SRPTensor to a torch-saveable payload.

Return type:: dict[str, Any]

static from_dict(payload, *, validate=True)[source]

Deserialize SRPTensor from payload produced by to_dict().

Return type:

SRPTensor

Parameters:

payload (dict[str, Any])
validate (bool)

static from_dense(x, k, *, score_mode='abs')[source]

Project dense tensor to fixed-k SRP row-packed representation.

Accepts shape (*prefix, cols_total) and flattens prefix dims into rows. Stores signed values gathered from original x at selected indices.

Return type:

SRPTensor

Parameters:

x (Tensor)
k (int)
score_mode (Literal['abs', 'raw', 'relu'])

Sparse Autoencoders

class compresso.TopKSAE(input_dim, hidden_dim, k, tied=False, decoder_bias=False, pre_act=None, post_sparsify=None, encoder=None, decoder=None, sparsify_score_mode='abs', sparsify_ste_alpha=0.0)[source]

Sparse autoencoder with hard top-k bottleneck.

Parameters:

input_dim (int) – Dimensionality of the input (and reconstruction).
hidden_dim (int) – Width of the sparse code layer.
k (int) – Number of active features per sample.
tied (bool) – If True, decoder weight is the transpose of the encoder weight.
pre_act (Optional[Module]) – Optional activation applied to encoder output before sparsification.
decoder_bias (bool)
post_sparsify (Optional[nn.Module])
encoder (Optional[nn.Module])
decoder (Optional[nn.Module])
sparsify_score_mode (str)
sparsify_ste_alpha (float)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

get_decoder_weight()[source]

Return the effective decoder weight matrix (input_dim, hidden_dim).

Return type:: Tensor

forward(x)[source]

Returns:

reconstruction (Tensor (B, input_dim))
codes (Tensor (B, hidden_dim) — sparse, exactly k nonzeros per row)
stats (dict[str, Tensor])

Parameters:

x (Tensor)

class compresso.TopKSAEConfig(hidden_dim=4096, k=128, decoder_bias=False, pre_act=None, post_sparsify=None, encoder=None, decoder=None, sparsify_score_mode='abs', sparsify_ste_alpha=0.01, noise_type='none', noise_scale='global_rms', noise_level=0.1, standard_scaler_mean=False, standard_scaler_scale='none', standard_scaler_loss_space='original', alpha_loss=0.01, l1_penalty=0.0, batch_size=128, shuffle=True, seed=42, epochs=10, validation_frac=None, patience=None, min_delta=0.0, restore_best_weights=True, lr=0.001, weight_decay=0.0, decay=False, compile=False, device='cpu', show_progress=True, srp_score_mode='abs', log_prefix='TopKSAE', log_every_n_steps=0)[source]

Configuration for TopKSAETrainer.

Parameters:

hidden_dim (int) – Width of the SAE code layer.
k (int) – Number of active code features per row.
decoder_bias (bool) – Whether the default decoder linear layer uses a bias.
pre_act (Module | None) – Optional module applied to encoder output before sparsification.
post_sparsify (Module | None) – Optional module applied to sparse codes after top-k. For example, L1Normalize().
encoder (Module | None) – Optional custom modules. When omitted, TopKSAE uses linear encoder and decoder layers.
decoder (Module | None) – Optional custom modules. When omitted, TopKSAE uses linear encoder and decoder layers.
sparsify_score_mode (Literal['abs', 'raw', 'relu']) – Top-k scoring mode: "abs", "raw", or "relu".
sparsify_ste_alpha (float) – Straight-through estimator leakage for non-selected positions.
noise_type (Literal['none', 'gaussian']) – Optional corruption applied to training inputs. "none" leaves inputs unchanged and "gaussian" adds Gaussian noise.
noise_scale (Literal['absolute', 'global_rms', 'feature_std']) – Scaling used for Gaussian noise. "absolute" uses embedding coordinate units, "global_rms" uses one training-set-derived scale, and "feature_std" scales each input feature separately.
noise_level (float) – Gaussian standard deviation for absolute scaling, or a dimensionless multiplier for adaptive scaling.
standard_scaler_mean (bool) – Subtract the per-feature training mean from inputs before the SAE, and add it back on its reconstruction. Off by default.
standard_scaler_scale (Literal['none', 'feature_std', 'global_rms']) –
Division applied after centering, and undone on the reconstruction. "none" leaves magnitudes alone. "feature_std" divides each feature by its own standard deviation, matching sklearn.preprocessing.StandardScaler; it gives every feature unit variance, which flattens the relative importance of coordinates and bends the embedding geometry, so adaptive noise_scale is rejected alongside it. "global_rms" divides everything by one scalar, the root mean per-feature variance: coordinates land at unit scale for the encoder while every angle and every distance ratio is preserved exactly, since a uniform scale is not a distortion.

Statistics are fitted on the training rows with correction=0, and constant features keep a scale of 1.

Neither scaling mode sits well with a normalizing post_sparsify: unit-norm codes carry no magnitude, so the rescale has to be undone by the decoder alone and converges several times worse. That combination warns rather than raising, since it is merely a bad trade rather than a contradiction. standard_scaler_mean is unaffected, because centering barely moves the magnitude.
standard_scaler_loss_space (Literal['original', 'scaled']) – Space the reconstruction loss is measured in when standard scaling is active. "original" un-scales the reconstruction and compares it to the raw input, keeping the objective and reported metrics identical to an unscaled run. "scaled" compares in standardized space, which weights every feature equally instead of by its variance.
alpha_loss (float) – Mixture weight for cosine loss. Training loss is alpha_loss * (1 - cosine_similarity) + (1 - alpha_loss) * mse.
l1_penalty (float) – Optional penalty on mean absolute sparse code activation.
batch_size (int) – Number of embedding rows per training batch.
shuffle (bool) – Whether to shuffle training rows between epochs.
seed (int) – Random seed used for row shuffling and Torch initialization.
epochs (int) – Maximum number of training epochs. Early stopping can end training before this many epochs have run.
validation_frac (float | None) – Optional fraction of input rows held out for validation. Mutually exclusive with the validation_embeddings argument of TopKSAETrainer.fit(). When both are unset, no validation pass runs and early stopping is unavailable.
patience (int | None) – Number of consecutive epochs without a validation improvement tolerated before training stops. None disables early stopping. Requires a validation set.
min_delta (float) – Smallest decrease in validation loss that counts as an improvement.
restore_best_weights (bool) – If True, reload the weights of the best-scoring epoch once training ends. Only applies when validation is active.
lr (float) – AdamW optimizer parameters.
weight_decay (float) – AdamW optimizer parameters.
decay (bool) – If True, use cosine learning-rate decay from lr to zero across the configured training epochs.
compile (bool) – If True, call torch.compile on the SAE when available.
device (str | device) – Device used for training and transforms.
show_progress (bool) – Whether to show a tqdm progress bar when tqdm is installed. Ignored whenever a logger is in play, since a bar and a log stream would report the same numbers twice; that rule is absolute, so a per-call show_progress=True does not reinstate the bar alongside a logger. Individual calls can override this default.
srp_score_mode (Literal['abs', 'raw', 'relu']) – Score mode used by SRPTensor.from_dense during transform.
log_prefix (str) – Bracketed tag put in front of every logged line, so one job’s output stays greppable when several share a log stream.
log_every_n_steps (int) – With a logger, also log every N-th batch of an epoch or of an inference pass. 0 logs epoch and pass boundaries only, which is enough unless a single epoch or pass runs for minutes.

TopKSAETrainer builds on the sparse embedding compression method described in Sparse embedding compression and TopKSAETrainer.

class compresso.TopKSAETrainer(config=None, logger=None)[source]

Efficient fit/transform wrapper around compresso.TopKSAE.

The trainer is intended for dense embedding matrices, such as item embeddings from a recommender or semantic embeddings from a text encoder. It exposes a compact sklearn-like API:

>>> trainer = TopKSAETrainer(TopKSAEConfig(k=32, epochs=300))
>>> trainer.fit(embeddings)
>>> sparse = trainer.transform(embeddings)

transform returns an SRPTensor containing sparse codes. Use reconstruct if dense reconstructions are needed.

Parameters:

config (TopKSAEConfig | None) – Trainer configuration. Defaults to TopKSAEConfig.
logger (Any | None) –
Default sink for progress lines, any object with an info(str) method: a logging.Logger, a service’s own logger, or a shim around print. Duck-typed on purpose, so compresso needs no logging dependency and callers keep their own. When given, runs report themselves line by line instead of through tqdm, which suits a container with no tty; when omitted, nothing about the trainer changes.

fit, fit_transform, encode, reconstruct, and transform each take logger and show_progress of their own, overriding this one for a single call. A sink describes the job that is running rather than the model, so it is never saved in state_dict() and never carried by from_state_dict() — which is also what keeps a checkpoint picklable when the sink holds a socket or an HTTP session.

property is_built: bool: Whether the underlying TopKSAE model has been initialized.

build(input_dim)[source]

Initialize model and optimizer for inputs of size input_dim.

Return type:: TopKSAETrainer
Parameters:: input_dim (int)

to(device)[source]

Move the underlying model to device and return self.

Return type:: TopKSAETrainer
Parameters:: device (str | device)

train_step(batch)[source]

Run one optimization step and return detached training stats.

Return type:: dict[str, Tensor]
Parameters:: batch (Tensor)

eval_step(batch)[source]

Score one batch without corruption and return detached stats.

Validation inputs are never passed through _corrupt(), so the reported loss is free of the per-epoch noise draw that makes training loss a poor early-stopping signal for denoising configurations.

Return type:: dict[str, Tensor]
Parameters:: batch (Tensor)

fit(embeddings, *, validation_embeddings=None, logger=<inherit>, show_progress=<inherit>)[source]

Train the SAE on dense embeddings and return self.

Parameters:

embeddings (ndarray | Tensor) – Training rows. When config.validation_frac is set, the held-out part is split off from these rows before any training statistics, including adaptive noise scales, are fitted.
validation_embeddings (ndarray | Tensor | None) – Explicit validation rows, for when the split is made by the caller. Mutually exclusive with config.validation_frac.
logger (Any | None) – Reporting sink for this call only, overriding the one given to the constructor. None disables log lines and drops an inherited progress bar unless show_progress is passed explicitly.
show_progress (bool | _Inherit) – tqdm bar for this call only, overriding config.show_progress. A logger suppresses the bar either way.

Return type:

TopKSAETrainer

encode(embeddings, *, logger=<inherit>, show_progress=<inherit>)[source]

Return dense sparse-code tensor produced by the trained SAE.

Parameters:

embeddings (ndarray | Tensor) – Rows to run through the fitted SAE.
logger (Any | None) – Reporting sink for this call only, overriding the one given to the constructor. None disables log lines and drops an inherited progress bar unless show_progress is passed explicitly.
show_progress (bool | _Inherit) – tqdm bar for this call only, overriding config.show_progress. A logger suppresses the bar either way.

Return type:

Tensor

reconstruct(embeddings, *, logger=<inherit>, show_progress=<inherit>)[source]

Return dense reconstructions for embeddings.

Parameters:

embeddings (ndarray | Tensor) – Rows to run through the fitted SAE.
logger (Any | None) – Reporting sink for this call only, overriding the one given to the constructor. None disables log lines and drops an inherited progress bar unless show_progress is passed explicitly.
show_progress (bool | _Inherit) – tqdm bar for this call only, overriding config.show_progress. A logger suppresses the bar either way.

Return type:

Tensor

transform(embeddings, *, logger=<inherit>, show_progress=<inherit>)[source]

Encode embeddings and return sparse codes as an SRPTensor.

Each batch is packed as it is produced, so the dense (n, hidden_dim) code matrix is never held whole: peak memory is the packed (n, k) result plus one batch. Top-k runs per row, so the result is identical to packing the full dense matrix in one go.

logger and show_progress override the trainer’s own reporting for this call, as on encode(). In particular, logger=None disables log lines and drops an inherited progress bar unless show_progress is passed explicitly.

Return type:

SRPTensor

Parameters:

embeddings (ndarray | Tensor)
logger (Any | None)
show_progress (bool | _Inherit)

fit_transform(embeddings, *, validation_embeddings=None, logger=<inherit>, show_progress=<inherit>)[source]

Fit the SAE and return encoded sparse codes as an SRPTensor.

Codes are returned for every row of embeddings, including any rows held out for validation by config.validation_frac.

logger and show_progress cover both phases, so a long transform is reported by whatever reported the fit. In particular, logger=None disables log lines and drops an inherited progress bar unless show_progress is passed explicitly.

Return type:

SRPTensor

Parameters:

embeddings (ndarray | Tensor)
validation_embeddings (ndarray | Tensor | None)
logger (Any | None)
show_progress (bool | _Inherit)

state_dict()[source]

Return a saveable trainer state dictionary.

Return type:: dict[str, Any]

load_state_dict(state, *, load_optimizer=True)[source]

Restore trainer state, including fitted denoising statistics.

Return type:

TopKSAETrainer

Parameters:

state (dict[str, Any])
load_optimizer (bool)

classmethod from_state_dict(state, *, device=None, load_optimizer=True)[source]

Construct a trainer from state_dict() output.

Return type:

TopKSAETrainer

Parameters:

state (dict[str, Any])
device (str | device | None)
load_optimizer (bool)

class compresso.L1Normalize(*args, **kwargs)[source]

Apply row-wise L1 normalization.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

args (Any)
kwargs (Any)

forward(x)[source]

Normalize each row by its L1 norm.

Parameters:: x (Tensor) – Input tensor whose last dimension is normalized.
Returns:: Tensor with the same shape as x and unit L1 norm along the last dimension where possible.
Return type:: Tensor

class compresso.L2Normalize(*args, **kwargs)[source]

Apply row-wise L2 normalization.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

args (Any)
kwargs (Any)

forward(x)[source]

Normalize each row by its L2 norm.

Parameters:: x (Tensor) – Input tensor whose last dimension is normalized.
Returns:: Tensor with the same shape as x and unit L2 norm along the last dimension where possible.
Return type:: Tensor

Sparse Parameters

class compresso.MaskedParam(weight, k_target, k_schedule=None, num_stages=10, stability_window=5, change_threshold=0.01, sparsity='row', allow_regrowth=True, score_mode='abs', ste_alpha=1.0, post_norm_l1=False)[source]

Dense parameter with a binary mask and an internal pruning schedule.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

weight (Tensor)
k_target (int)
k_schedule (Sequence[int] | None)
num_stages (int)
stability_window (int)
change_threshold (float)
sparsity (Literal['row', 'col'])
score_mode (Literal['abs', 'raw', 'relu'])
ste_alpha (float)
post_norm_l1 (bool)

forward()[source]

Return the current masked or top-k-projected weight tensor.

When the mask is frozen, the stored binary mask is applied directly to weight. Otherwise, a differentiable top-k projection is computed using the current pruning stage and straight-through estimator settings.

Returns:: Dense weight tensor with inactive entries set to zero.
Return type:: torch.Tensor

select_rows(row_indices)[source]

Alias for row indexing without full masked-weight materialization.

Return type:: Tensor
Parameters:: row_indices (slice | Tensor | Sequence[int])

step_mask()[source]: Compute mask and add history

rewind()[source]: Back to original init with new k

maskedparam_to_coo()[source]: Export packed fixed-k COO according to self.dim: - dim==1 -> row-packed, fixed k per row - dim==0 -> col-packed, fixed k per col

to_srp_param()[source]

Export the exact final or frozen row mask as an SRPParam.

Return type:: SRPParam

maskedparam_to_srp()[source]

Compatibility alias for to_srp_param().

Return type:: SRPParam

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

class compresso.SRPParam(cols, values, shape, *, validate=True)[source]

Structured row-packed sparse parameter.

Represents a matrix A of shape (rows, cols_total) stored as fixed-k nonzeros per row.

Storage:

cols: (rows, k) int64 buffer with fixed indices per row.
values: (rows, k) trainable parameter.

Semantics:

A[r, cols[r, j]] += values[r, j]

Duplicates within a row accumulate with scatter-add semantics. This format is useful for row-packed sparse matrix multiplication against a dense matrix B with shape (cols_total, out).

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

cols (Tensor)
values (Tensor)
shape (Tuple[int, int])
validate (bool)

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:: str

build_coo(*, dtype=None, coalesce=True)[source]

Convert to a torch sparse COO tensor (rows, cols_total). Duplicates are kept; coalesce() will sum them.

Return type:

Tensor

Parameters:

dtype (dtype | None)
coalesce (bool)

select_rows(row_indices)[source]

Alias for gradient-preserving row indexing.

Return type:: SRPTensor
Parameters:: row_indices (slice | Tensor | Sequence[int])

static from_dense(A, k, *, mode='topk_abs', allow_duplicates=False)[source]

Build SRPParam from a dense matrix.

mode:

"topk_abs": pick top-k absolute values per row and store signed values.
"random_k": pick k random columns per row and store those values.

Note: this is a projection of dense -> fixed-k sparse.

Return type:

SRPParam

Parameters:

A (Tensor)
k (int)
mode (Literal['topk_abs', 'random_k'])
allow_duplicates (bool)

static from_sparse_coo(sp, *, k=None, require_row_packed_fixed_k=True)[source]

Convert from a COO sparse tensor to SRPParam.

Return type:

SRPParam

Parameters:

sp (Tensor)
k (int | None)
require_row_packed_fixed_k (bool)

Cases:

If require_row_packed_fixed_k=True:
expects exactly k nnz per row (same k for all rows), and we will pack rows. If k is None, we infer it. If not possible -> error.
If require_row_packed_fixed_k=False:
we will project each row to top-k by abs value (needs k provided).

forward()[source]

Return the sparse representation as an SRPTensor.

Returns:: Row-packed sparse representation backed by this parameter’s column indices and trainable values.
Return type:: SRPTensor

Utilities

class compresso.SparsityController(model, mask_update_interval=10, freeze_at_schedule_end=True, method='all')[source]

Global dispatcher for MaskedParam-based pruning.

Responsibilities:

Keep track of global step.
Every mask_update_interval steps, call step_mask() on all MaskedParams.
If any MaskedParam’s stage is completed, call rewind() on all of them and signal that the optimizer should be reset.
Once all MaskedParams have schedule_done=True, optionally freeze masks and stop scheduling.

Parameters:

model (Module)
mask_update_interval (int)
freeze_at_schedule_end (bool)

step()[source]

Call this once per optimizer step.

Returns:

‘rewind_triggered’: bool, whether a global rewind happened
’all_schedules_done’: bool, whether all MaskedParams finished schedule

Return type:

dict

compresso.exponential_decay(n_init, n_final, k)[source]: Smooth exponential decay from n_init to n_final across k steps, guaranteed never to go below n_final.

compresso.save_srp_tensor(path, srp)[source]

Save SRPTensor payload with torch.save.

Return type:

None

Parameters:

path (str | Path)
srp (SRPTensor)

compresso.load_srp_tensor(path, map_location=None, *, validate=True)[source]

Load SRPTensor payload produced by save_srp_tensor().

Return type:

SRPTensor

Parameters:

path (str | Path)
validate (bool)