peal.sparse_dictionaries.batch_topk_network.BatchTopKNetwork

class peal.sparse_dictionaries.batch_topk_network.BatchTopKNetwork

Bases: SparseAutoencoderBase

Batch top-K sparse autoencoder.

During training the top_k * batch_size largest codes of the whole batch survive; at evaluation a running threshold stands in for that pooled selection, since a single sample has no batch to pool over.

Methods

compute_activations	Return (dense_relu_codes, sparse_codes) for a preprocessed x.
decode	Reconstruct from sparse codes, undoing the cached standardisation.
encode	Sparse codes of x; the standardisation is cached for decode().
forward	Full training step output: reconstruction plus the loss breakdown.
get_auxiliary_loss	Ask the dead atoms to explain the current residual, so they revive.
get_loss_dict	Reconstruction, sparsity, dead-atom revival and orthogonality terms.
get_orthogonality_loss	Penalise off-diagonal atom correlation; PEAL's own addition, off by default (orth_coeff 0).
update_threshold

__init__(cfg)

compute_activations(x)

Return (dense_relu_codes, sparse_codes) for a preprocessed x.

forward(x)

Full training step output: reconstruction plus the loss breakdown.

encode(x)

Sparse codes of x; the standardisation is cached for decode().

decode(acts_topk)

Reconstruct from sparse codes, undoing the cached standardisation.

get_loss_dict(x, x_reconstruct, acts, acts_topk, x_mean, x_std)

Reconstruction, sparsity, dead-atom revival and orthogonality terms.

get_auxiliary_loss(x, x_reconstruct, acts)

Ask the dead atoms to explain the current residual, so they revive.

get_orthogonality_loss()

Penalise off-diagonal atom correlation; PEAL’s own addition, off by default (orth_coeff 0).

update_threshold(acts_topk, lr=0.01)

Track the smallest surviving code, the inference-time cut-off.