structboost.BAEConfig

class structboost.BAEConfig(latent_dim=10, decoder_hidden_dims=(64,), decoder_activation='tanh', split_softmax=False, boosting_stepno=50, boosting_nu=0.1, boosting_csf=0.9, boosting_independent=True, boosting_precompute_covcache='auto', nuisance_ridge=0.0, prior_mode='frozen', disentanglement='orthogonal', disentanglement_alpha=1.0, disentanglement_lambda=0.01, target_optim_lr=1.0, decoder_lr=0.001, decoder_weight_decay=0.0, max_iterations=1000, enable_early_stopping=False, early_stopping_patience=50, batch_size=512, seed=None, device='cpu', diagnostics=False)

Bases: object

Configuration for Boosting Autoencoder.

The revised BAE architecture uses a linear encoder (optimized via boosting) and an MLP decoder (optimized via SGD).

Parameters:

• latent_dim (int)

• decoder_hidden_dims (tuple[int, ...])

• decoder_activation (Literal['tanh', 'relu', 'leaky_relu', 'elu'])

• split_softmax (bool)

• boosting_stepno (int)

• boosting_nu (float)

• boosting_csf (float)

• boosting_independent (bool)

• boosting_precompute_covcache (bool | Literal['auto'])

• nuisance_ridge (float)

• prior_mode (Literal['frozen', 'anchored'])

• disentanglement (Literal['none', 'correlation', 'orthogonal'])

• disentanglement_alpha (float)

• disentanglement_lambda (float)

• target_optim_lr (float)

• decoder_lr (float)

• decoder_weight_decay (float)

• max_iterations (int)

• enable_early_stopping (bool)

• early_stopping_patience (int)

• batch_size (int)

• seed (int | None)

• device (str | Any)

• diagnostics (bool)

latent_dim

Number of latent dimensions.

Type:

int

decoder_hidden_dims

Hidden layer sizes for the decoder MLP, in forward order from the latent code to the gene output. Default (64,) — a single hidden layer.

On simulated data with planted gene programs, one hidden layer of 64 units recovered markers as well as or better than (128,) (marker-recovery F1 0.99–1.00 across four scenarios spanning 500–3000 genes), reconstructed to ~91–95% of the linear ceiling, and trained faster. A linear decoder (()) is a valid faster, more interpretable choice but loses recall on harder data (F1 down to 0.83, reconstruction 65–73% of ceiling). Deeper or wider decoders did not help and often hurt marker recovery: a funnel such as (128, 64) dropped to F1 0.80 because a narrow layer immediately before the gene output distorts the reconstruction gradient that the boosting target is built from — a BAE-specific effect, since that gradient is what the encoder is fitted against. If more capacity is needed, widen ((64, 128)) rather than narrowing toward the output. These are simulation results; confirm on real data before assuming (64,) has enough capacity for genuinely nonlinear expression structure.

Type:

tuple[int, …]

decoder_activation

Activation function for decoder hidden layers. "tanh" (default) is the most stable; "elu" matched it; "relu" was unstable at larger widths (marker-recovery F1 collapsing to 0.80 at 256 units).

Type:

Literal[‘tanh’, ‘relu’, ‘leaky_relu’, ‘elu’]

split_softmax

If True, apply split-softmax transformation between encoder and decoder. Each latent dimension z_i is paired with −z_i (interleaved) and softmax- normalized onto the 2d-dimensional simplex: σ_split(z) = softmax((z_1, −z_1, …, z_d, −z_d)). The decoder then receives a 2d-dimensional compositional input, enabling soft clustering of cells into 2d groups. Default is False (standard unconstrained latent space).

Type:

bool

boosting_stepno

Number of boosting iterations per training iteration.

Type:

int

boosting_nu

Boosting learning rate.

Type:

float

boosting_csf

Cumulative shrinkage factor for boosting.

Type:

float

boosting_independent

If True, reset boosting state for each latent dimension (recommended).

Type:

bool

nuisance_ridge

Ridge strength for the batch covariates when batch_integration_mode includes “encoder”, expressed relative to each encoded column’s squared norm. Mandatory gene coefficients are never penalized. Default 0 preserves exact least squares.

This is a numerical-stability knob rather than a modelling one. Boosting refits the mandatory block jointly at every step, and that block is the mandatory genes together with the batch columns. When two of those are collinear the solve has no unique answer and raises; when they are merely close to collinear it returns a large, sign-unstable answer and raises nothing. Reaching for a small value here, 1e-3 say, is the remedy in the second case. It is never applied automatically, because ridge changes the estimates and doing so silently would fit a different model than the one asked for.

Type:

float

prior_mode

How the prior encoder weight matrix is treated when the model was built by BAE.from_reference(); ignored otherwise. "frozen" (default) never boosts the prior columns, so they remain bitwise equal to the reference matrix and the transferred gene programs are literally the reference’s. "anchored" boosts them too, but restarts each iteration from the fixed original matrix rather than from the previous iteration’s result — boosting from an offset model, so deviation is bounded by boosting_stepno and never accumulates. Unanchored re-fitting is not offered: a prior used merely as a starting point is forgotten within a few hundred iterations, since the encoder support random-walks and two runs end no more similar than two unrelated ones.

Type:

Literal[‘frozen’, ‘anchored’]

boosting_precompute_covcache

Whether to compute the full p×p covariance matrix before training. "auto" (default) does so whenever the 8·p² byte matrix fits a conservative share of system memory; True and False are honoured exactly. The resolved decision is recorded in adata.uns["bae"]["boosting_precompute_covcache"].

This is a speed setting first and a memory setting second, which is not obvious. The alternative is a dict-backed column cache: a covariance column is computed the first time its feature is selected and reused across targets and iterations, so memory scales with the number of distinct selected features (8·p bytes per column) rather than with p². That sounds strictly cheaper, and in memory it is — but building all p columns at once is a single compute-bound matrix product running near hardware peak, while fetching them one at a time is a sequence of memory-bound matrix-vector products. Measured on real data, precomputing wins from roughly p/59 distinct selected features onwards, and a fit passes that within its first iteration, where up to boosting_stepno × latent_dim distinct features can enter. End-to-end it measured 1.7–2.3× faster at p=3,000–10,000.

The cost is memory, and it is real: 8·p² is 32 MB at p=2,000, 800 MB at p=10,000 and 3.2 GB at p=20,000, transiently doubled while the matrix is built. That is what "auto" guards against; set False explicitly on a memory-constrained machine.

Type:

bool | Literal[‘auto’]

disentanglement

Exploratory, under active development.

Latent-dimension disentanglement method, "orthogonal" by default. That method replaces the boosting targets with the nearest mutually orthogonal set of the same column norms — symmetric Löwdin orthogonalization, computed from the thin SVD. "none" applies no constraint.

Its strength is set by disentanglement_alpha: full orthogonality at the default 1.0. Decorrelation can be softened by choosing an alpha strictly between 0 and 1, which applies the constraint only partially.

"correlation" is the alternative, adding a soft squared-correlation penalty to the target objective rather than transforming the targets. It requires manual tuning — its effect is negligible at small lambda, and larger values give stronger decorrelation — whereas "orthogonal" has a bounded dial that needs no calibration.

Expect either method to cost some biological structure: decorrelation is an extra constraint on the latent space and real gene programs are not orthogonal — overlapping pathways and shared markers are the norm — so it trades fidelity to that structure, and the interpretability that rests on it, for a less redundant representation. Measured against ground truth on simulated data, orthogonalization moved marker-recovery F1 from 0.98 to 0.88, and on Tasic and human pancreas it lowered the share of latent variance explained by the annotated cell type. "none" therefore remains a reasonable choice.

Type:

Literal[‘none’, ‘correlation’, ‘orthogonal’]

disentanglement_alpha

Strength of the orthogonality constraint, used only when disentanglement="orthogonal". The targets become (1 - alpha) * targets + alpha * orthogonalized, so 1.0 (default) makes them exactly orthogonal and 0.0 leaves them untouched, which is equivalent to disentanglement="none".

Lower it to soften the constraint when full orthogonality costs more biological structure than the redundancy it removes is worth. Unlike disentanglement_lambda the scale is bounded and interpretable, so no calibration sweep is needed to find a usable range — the endpoints are known and the response is monotone in between.

Type:

float

disentanglement_lambda

Strength of the soft correlation penalty; used only when disentanglement="correlation". Larger values decorrelate more strongly. The penalty is normalized over latent-dimension pairs and scaled with the number of cells, so the value means the same thing at different dataset sizes.

Raised to 1e-2 in 0.5.0. The previous 1e-4 was measured to be inert — on simulated data, Tasic mouse cortex and human pancreas it moved the mean absolute latent correlation by less than its own seed-to-seed noise — and so was 1e-3. The penalty begins to act at 1e-2 and is strong at 1e-1. On pancreas, mean |corr| ran 0.103 (off), 0.100 (1e-4), 0.112 (1e-3), 0.080 (1e-2), 0.061 (1e-1), the last costing noticeably more cell-type structure.

This is a starting point rather than a tuned optimum, and the right value is dataset-dependent. Sweep upward — 1e-2, 3e-2, 1e-1, 3e-1 — watching latent correlation against whatever downstream structure matters to you.

Type:

float

target_optim_lr

Step size for computing boosting targets via gradient descent on z. Targets are computed as z* = z - lr * ∂L_target/∂z (single gradient step), where L_target sums the squared error over cells and averages it over genes. That convention — rather than the elementwise mean used for the decoder update and every reported loss — makes a cell’s target step independent of how many cells the dataset contains, so a given target_optim_lr means the same thing at every dataset size.

Leave this at 1.0. It is exposed because it is a real coefficient of the method, not because it is a tuning knob: it scales the whole coupling between the two optimizers, and the quantity it scales has no natural unit — ∂L/∂z tracks the decoder’s magnitude, which grows by orders of magnitude over a fit. A value that suits one training stage will not suit the next, so the alternation is what adapts, and this stays fixed.

Note

An earlier formulation took the target gradient from the elementwise mean MSE, making it a factor n_cells smaller. A value ported from such a setup must be divided by the number of cells.

Type:

float

decoder_lr

Learning rate for decoder SGD optimization.

Type:

float

decoder_weight_decay

Weight decay (L2 penalty) for the AdamW decoder optimizer. Default is 0.0 (no penalty). Increase for decoder weight regularisation (typical range: 1e-5 to 1e-2).

Type:

float

max_iterations

Maximum number of training iterations.

Type:

int

enable_early_stopping

If True, stop training early when the checkpoint-selection loss does not decrease for early_stopping_patience iterations. This is the decoder training MSE for disentanglement="none" and "orthogonal"; for "correlation" it additionally includes the weighted disentanglement penalty. Default False: train for exactly max_iterations.

The default changed in 0.4.0 because the criterion is a convergence check being used as a quality check, and it stops far too early. There is no validation split, so the training loss cannot see a model that is beginning to memorize; measured against ground truth it stopped at iteration 90 on one dataset whose marker recovery peaked at 259, and at 196 on a simulated scenario peaking at 560 (returning F1 0.502 against 0.787). Across three real datasets the best iteration ranged from 154 to 1975, always past where patience fires.

No stopping rule replaced it, because none was found that works. Latent stability, support overlap and held-out reconstruction were each measured as candidates; the representation settles long before gene selection does, and nothing observable tracks the quality peak. Run the iteration budget and let BAE.stability_selection() absorb the variance instead — that is what it is for.

Type:

bool

early_stopping_patience

Number of consecutive iterations without checkpoint-selection loss improvement required to trigger early stopping. Only used if enable_early_stopping is True. Default is 50.

Type:

int

batch_size

Minibatch size for decoder SGD updates, and with it the number of decoder updates each training iteration performs: the cells are shuffled and partitioned, so one iteration runs ceil(n_cells / batch_size) steps and every cell contributes to exactly one of them.

That coupling is the point, and it replaced a separate decoder_updates_per_iteration in 0.5.0. A fixed step count made a cell’s participation depend on dataset size — at 10 steps of 512 the decoder saw every cell below 5,120 cells, 31% at 16,000 and 5% at 100,000 — while the boosting half of the alternation is full-batch at every size: z* is computed on all cells and the encoder is re-solved on all cells, every iteration. Tying the decoder budget to the data makes both halves consume the same cells per iteration.

The cost is that fit time now grows with n_cells, since max_iterations does not decay to compensate. Raising batch_size buys the time back at the price of fewer, noisier steps.

Type:

int

seed

Random seed for reproducibility. If None, no seed is set. Controls train/val split, weight initialization, and SGD shuffling.

Type:

int | None

device

Device for computation (‘cpu’, ‘cuda’, or torch.device).

Type:

str | Any

diagnostics

If True, collect per-iteration training diagnostics into a TrainingReport (see BAE.training_report) and additionally show the relative encoder weight change (dW) and the number of selected genes (n_sel) in the progress bar. This adds full-data forward and backward passes per iteration, so it is off by default and diagnostics=False costs exactly what training cost before. Diagnostics are read-only: enabling them does not change the fitted model.

Type:

bool

Methods

__init__([latent_dim, decoder_hidden_dims, ...])

Attributes

batch_size
boosting_csf
boosting_independent
boosting_nu
boosting_precompute_covcache
boosting_stepno
decoder_activation
decoder_hidden_dims
decoder_lr
decoder_weight_decay
device
diagnostics
disentanglement
disentanglement_alpha
disentanglement_lambda
early_stopping_patience
enable_early_stopping
latent_dim
max_iterations
nuisance_ridge
prior_mode
seed
split_softmax
target_optim_lr