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)

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, alpha_loss=0.01, l1_penalty=0.0, batch_size=128, shuffle=True, seed=42, epochs=10, lr=0.001, weight_decay=0.0, decay=False, compile=False, device='cpu', show_progress=True, srp_score_mode='abs')[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.
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) – Number of training epochs.
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.
srp_score_mode (Literal['abs', 'raw', 'relu']) – Score mode used by SRPTensor.from_dense during transform.

class compresso.TopKSAETrainer(config=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)

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)

fit(embeddings)[source]

Train the SAE on dense embeddings and return self.

Return type:: TopKSAETrainer
Parameters:: embeddings (ndarray | Tensor)

encode(embeddings)[source]

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

Return type:: Tensor
Parameters:: embeddings (ndarray | Tensor)

reconstruct(embeddings)[source]

Return dense reconstructions for embeddings.

Return type:: Tensor
Parameters:: embeddings (ndarray | Tensor)

transform(embeddings)[source]

Encode embeddings and return sparse codes as an SRPTensor.

Return type:: SRPTensor
Parameters:: embeddings (ndarray | Tensor)

fit_transform(embeddings)[source]

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

Return type:: SRPTensor
Parameters:: embeddings (ndarray | Tensor)

state_dict()[source]

Return a saveable trainer state dictionary.

Return type:: dict[str, Any]

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

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

maskedparam_to_srp()[source]

Export fixed-k row-packed SRPParam.

Only valid for row-wise sparsity (self.dim == 1), because SRPParam stores exactly k entries per ROW.

Uses:

frozen mask if mask_frozen
otherwise current-stage top-k mask, k_current

Returns:

SRPParam with cols shaped (rows, k_current), values shaped (rows, k_current), and shape (rows, cols).

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]

Fast row selection, returns a NEW SRPParam with rows=len(row_indices), same cols_total and k. Reindexes rows to [0..R-1], structure is still row-packed.

Return type:: SRPParam
Parameters:: row_indices (Tensor)

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)