Fitting a model¶
The minimal path¶
from structboost import BAE, BAEConfig
config = BAEConfig(latent_dim=10)
model = BAE(adata.n_vars, config)
model.fit(adata) # adata.X must be z-scored
latent = adata.obsm["X_bae"]
weights = adata.varm["BAE_encoder_weights"]
fit returns the model, so it chains. fit_transform(adata) is exactly
fit(...) followed by returning adata.obsm["X_bae"], with no extra behaviour, in
case you were looking for some.
Reading from a layer instead of adata.X:
model.fit(adata, layer="scaled")
The layer is remembered for every later call. See What a fit writes into your AnnData.
Configuration¶
BAEConfig has 24 fields. Most have measured defaults you
should leave alone. These are the ones worth thinking about.
Architecture¶
Field |
Default |
Notes |
|---|---|---|
|
|
Number of gene programs to learn. |
|
|
One hidden layer. See below. |
|
|
Most stable. |
|
|
Compositional latent, for soft clustering. |
Warning
Do not narrow the decoder toward the output. A funnel such as (128, 64)
dropped marker-recovery F1 to 0.80 in benchmarks, against 0.99–1.00 for (64,).
A narrow layer immediately before the gene output distorts the reconstruction
gradient that the boosting target is built from, and that gradient is what the
encoder is fitted against, so the damage lands on gene selection specifically. If
you need capacity, widen: (64, 128).
A linear decoder (decoder_hidden_dims=()) is a legitimate faster, more
interpretable choice, at a cost: F1 down to 0.83 and 65–73% of the linear ceiling
on harder data.
These are simulation results. Confirm on your data before assuming (64,) has
enough capacity for genuinely nonlinear structure.
Boosting: where sparsity comes from¶
Field |
Default |
Notes |
|---|---|---|
|
|
The sparsity dial. Caps genes per dimension per iteration. |
|
|
Boosting learning rate. See below before raising it. |
|
|
|
|
|
Reset boosting state per latent dimension. |
|
|
Stabilizes near-collinear batch covariates. See Batches and unwanted covariates. |
These defaults replaced an earlier 0.3 / 100 pairing, which over-selected
badly: marker-recovery F1 of 0.72/0.97/0.57 across three
simulated scenarios, against 0.98/1.00/0.98 for the current defaults.
nu and stepno are not independent. The effective ridge degrees of freedom of
a boosting step is nu, so a fit’s per-iteration complexity budget is roughly
stepno * nu: 5 for the current defaults, 30 for the rejected pairing.
Warning
boosting_nu=0.3 at the default stepno=50 is not safe under the current
stopping rule. Measured on the same three scenarios, it reaches a marker-recovery
F1 of 0.975–0.987 within 18–57 iterations — as good as 0.1, three to four times
faster — and then degrades to 0.57–0.78 if training continues, while the
training MSE keeps falling. Early stopping watches that MSE, so nothing detects it,
and on two of three seeds the fit ran to max_iterations.
On low-signal data the picture reverses: on the hard scenario (weak effects,
dropout, ambient contamination) 0.3 reaches F1 0.787 against 0.1’s 0.406, and
needs ~560 iterations to get there.
So the right value depends on the regime, and the stopping rule cannot currently
tell them apart. Until it can, 0.1 is the safe default. If you raise it, watch
n_iterations against max_iterations and variance_explained against
linear_ceiling() — a fit that exhausts its iteration budget
with variance explained above the linear ceiling is in the degrading regime.
boosting_precompute_covcache decides how the predictor Gram matrix is cached.
The default "auto" builds it in full up front whenever 8·p² bytes fit a
conservative share of system memory, and falls back to a lazy column cache when
they do not. True and False are honoured exactly, and the resolved decision
is recorded in adata.uns["bae"]["boosting_precompute_covcache"].
Note
This reads like a memory setting and is really a speed setting. The lazy
cache computes a column the first time its gene is selected, so its memory
scales with the number of distinct selected genes rather than with p² — it
is strictly cheaper in memory. But building all p columns at once is one
compute-bound matrix product running near hardware peak, while fetching them one
at a time is a sequence of memory-bound matrix-vector products. Precomputing
wins from roughly p/59 distinct selected genes onwards, and a fit passes that
in its first iteration, where up to boosting_stepno × latent_dim genes can
enter. Measured end-to-end: 1.6–1.9× faster at p = 2,000–8,000.
The cost is real memory: 8·p² is 32 MB at p = 2,000, 800 MB at 10,000 and
3.2 GB at 20,000, transiently doubled while it is built. That is what "auto"
guards. Set False explicitly on a memory-constrained machine.
Disentanglement¶
Exploratory
Under active development. Both methods are provisional, and the measurements below come from three datasets at three seeds each.
Latent dimensions can end up redundant. Since 0.5.0 the default is
disentanglement="orthogonal", which replaces the boosting targets with the
nearest mutually orthogonal set of the same column norms — symmetric Löwdin
orthogonalization, computed from the thin SVD as U @ Vt.
Field |
Default |
Notes |
|---|---|---|
|
|
|
|
|
Strength of the orthogonality constraint, in |
|
|
Correlation-penalty strength; larger decorrelates more. |
Softening it. disentanglement_alpha interpolates:
(1 - alpha) * targets + alpha * orthogonalized. At 1.0 the targets are exactly
orthogonal, at 0.0 untouched — so decorrelation can be softened by choosing an
alpha strictly between 0 and 1. The response is monotone and the scale is
bounded, so no calibration sweep is needed to find a usable range.
BAEConfig(latent_dim=10, disentanglement_alpha=0.5) # half-strength
BAEConfig(latent_dim=10, disentanglement="none") # off
"correlation" is the alternative, adding a soft squared-correlation penalty
to the target objective instead of transforming the targets, plus a variance
barrier so low correlation cannot be achieved by collapsing dimensions. It
requires manual tuning, and larger disentanglement_lambda gives stronger
decorrelation.
Warning
The old disentanglement_lambda=1e-4 default 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 did 1e-3. 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 default is now 1e-2; sweep upward from there
(1e-2, 3e-2, 1e-1, 3e-1) rather than downward.
Expect either method to cost 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 resting 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; on Tasic and pancreas it lowered the
share of latent variance the annotated cell type explains. "none" is a
reasonable choice, and disentanglement_alpha exists so the trade can be made
partially.
Optimization and stopping¶
Field |
Default |
|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
batch_size sets the decoder budget¶
Each iteration gives the decoder one shuffled pass over the cells:
ceil(n_cells / batch_size) AdamW steps, with every cell contributing to exactly
one. There is no separate step-count setting — decoder_updates_per_iteration
was removed in 0.5.0.
The reason is that the boosting half is full-batch at every dataset size: z* is
computed on all cells and the encoder is re-solved on all cells, every iteration.
A fixed step count made the decoder’s share shrink as data grew — at the old
10 × 512 that was every cell below 5,120, 31% at 16,000 and 5% at 100,000.
Warning
Two consequences worth planning for.
Fit time now grows with cell count, because max_iterations does not decay to
compensate. Raise batch_size to buy it back, at the price of fewer, noisier
steps.
On small datasets you may want a smaller batch_size than you think. At the
default 512, anything under ~1,000 cells gets one or two decoder steps per
iteration, and effects that need the decoder to move within an iteration weaken
or reverse. Measured on 300 cells, a PCA warm start raised the first
iteration’s loss at one step per iteration (1.083 against 1.023 for a zero start)
and only paid off once the pass held enough steps (0.920 against 1.000 at
batch_size=8). If a fit on few cells looks sluggish, lower batch_size before
reaching for anything else.
Leave target_optim_lr at 1.0. It is exposed because it is a real
coefficient of the method, not because it is a knob: it scales the entire
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, so a value that suits one training stage will not suit the next. The
alternation is what adapts; this stays fixed.
Early stopping is off by default¶
Changed in 0.4.0. 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 starting to memorize.
Measured against ground truth, patience-50 stopped at iteration 90 on a 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.
Warning
No stopping rule replaced it, because none was found that works. Latent stability, encoder-support overlap and held-out reconstruction were each measured as candidates. The representation settles long before gene selection does — on one dataset consecutive latent codes were rank-identical while the gene set still turned over 65% cumulatively — and no observable signal tracks the quality peak. The best iteration varies by an order of magnitude between datasets, so no fixed budget works either.
Run the iteration budget and let stability_selection()
absorb the variance. That is what it is for.
max_iterations=1000 is a defensible middle of the measured 154–1975 range, not an
optimum. Nothing better is available.
fit always restores the best checkpoint’s encoder, decoder and nuisance weights
at the end, whether or not early stopping triggered.
The objective early stopping watches is not always the training loss: with
disentanglement="correlation" it includes the weighted penalty, so that
restoring the best-reconstruction checkpoint cannot silently undo the constraint.
Per-call overrides beat the config, and are compared against None rather than
truthiness, so an explicit 0 overrides instead of falling back:
model.fit(adata, max_iterations=50, seed=0, verbose=False)
Warm starts¶
Exploratory
Under active development. Starting from an existing representation works, but it
is applied only once, it can silently change latent_dim, and its benefit is
conditional on the decoder having enough steps in the first iteration to follow
it — see the warning below.
Initialize the latent code from an existing embedding or a PCA instead of from zero:
model.fit(adata, init_pca=True)
model.fit(adata, init_obsm="X_pca")
model.fit(adata, init_pca=True, init_pretrain_epochs=20)
init_obsm and init_pca are mutually exclusive.
Warning
The warm start applies once, on the first iteration only. It sets the boosting targets. The encoder learns that representation, and the decoder is then trained against the encoder’s output.
It can change latent_dim silently-ish. If the supplied representation has a
different number of columns than config.latent_dim, the representation wins:
latent_dim is overwritten for this fit, the encoder is rebuilt, and a
UserWarning is emitted.
init_pretrain_epochs pre-fits the decoder against the fixed warm-start latent.
Measured trade-off: it lowers the initial loss substantially but does not
improve the converged loss, and large values noticeably reduce gene-selection
precision, because the decoder is tuned to a latent code the sparse encoder
cannot exactly reproduce. 20 is conservative.
None of these can be combined with a transfer model. See Transferring a model to a new dataset.
Reproducibility¶
BAEConfig(latent_dim=10, seed=42)
model.fit(adata, seed=42) # per-call override
A seeded fit is reproducible and does not disturb your own RNG state: the global torch stream is restored when fit returns, and NumPy’s
global generator is never touched at all. Constructing a BAE still advances the
torch generator, as constructing any nn.Module does.
Two unseeded fits still differ from each other, deliberately.