This walks the exact math executed by binn.fit(GCN(pathway_map, num_classes=C), data, ...)
(binn/train/loop.py) for a single epoch — forward pass through binn/nn/models.py's
GCN/Encoder, cross-entropy loss, backward pass through the pruning mask, and the
Adam parameter update. Everything below is read directly off the code, not idealized
textbook GCN math — including the parts that look unusual (double ReLU, decoupled-vs-coupled
weight decay) are called out where the implementation diverges from the textbook.
Notation: N cells (graph nodes), L pathway-hierarchy levels (so L-1 level transitions
plus one classifier head), layer widths units = [u_0, u_1, ..., u_{L-1}] where u_0 is the
gene count and u_{L-1} is the top pathway count. C = number of classes.
X ∈ R^{N × u_0}—data.x, gene-expression features per cell.edge_index→ adjacencyA ∈ {0,1}^{N × N}— the kNN cell-similarity graph (graphtools), not biological — this is the message-passing topology, independent of the pathway prior.y ∈ {0,...,C-1}^N—data.y, andtrain_mask ∈ {0,1}^Nselecting the training cells.- Binary pathway masks
M^{(i)} ∈ {0,1}^{u_{i+1} × u_i}fori = 0, ..., L-2, built once inbinn/pathways.py::masks_from_factorized:
M^{(i)}_{p, g} = 1 if gene/child-node g (level i) is an ancestor of pathway/parent p (level i+1)
= 0 otherwise
derived from the factorized pathway map by direct index assignment (mask_np[right, left] = 1.0) —
no learning involved, this is the fixed biological prior. A final all-ones mask
M^{(L-1)} ∈ {1}^{C × u_{L-1}} is appended for the classifier head (fully connected, unconstrained).
Each GCNConv's feature-transform weight is registered with
torch.nn.utils.prune.custom_from_mask(conv.lin, "weight", M^{(i)}). This does not store a
single weight tensor — it splits the parameter into a free tensor and a frozen mask, and
recomputes the effective weight on every forward call via a forward pre-hook:
W^{(i)}_orig ∈ R^{u_{i+1} × u_i} — the actual learnable nn.Parameter ("weight_orig")
M^{(i)} — fixed buffer, never updated by the optimizer ("weight_mask")
W^{(i)} = W^{(i)}_orig ⊙ M^{(i)} — effective weight used in forward, elementwise product
GCNConv keeps bias=False on conv.lin and, when self.bias=True (default False here),
holds a separate unmasked bias vector b^{(i)} ∈ R^{u_{i+1}} added post-aggregation. The
classifier head (nn.Linear) is a plain dense layer — its mask is all-ones so pruning is a
no-op there.
GCNConv (default add_self_loops=True, normalize=True, improved=False) computes symmetric-normalized
propagation over the self-looped adjacency:
 = A + I_N
D̂_{jj} = Σ_k Â_{jk} (degree of the self-looped graph)
Ŝ = D̂^{-1/2} Â D̂^{-1/2} (Ŝ ∈ R^{N×N}, fixed for the whole step — recomputed from edge_index, not cached)
For block i with input activations H^{(i)} ∈ R^{N × u_i} (H^{(0)} = X):
Z^{(i)} = H^{(i)} (W^{(i)})^T — feature transform, R^{N × u_{i+1}}
P^{(i)} = Ŝ Z^{(i)} — message passing / neighborhood aggregation
G^{(i)} = P^{(i)} + b^{(i)} (if bias, else P^{(i)}) — GCNConv raw output
Then Encoder.forward (binn/nn/models.py:280-290) applies, exactly as coded:
A1^{(i)} = ReLU(G^{(i)}) — ReLU after the conv (unconditional)
B1^{(i)} = BatchNorm1d(A1^{(i)}) — batch stats over the N node dimension
H^{(i+1)} = ReLU(B1^{(i)}) if i < L-1 (i.e. not the final block)
= B1^{(i)} if this is the last pathway block feeding straight into the classifier head shape check
H^{(i+1)} = Dropout_p(H^{(i+1)}, training=True) — only if i < L-1 and dropout > 0
Implementation quirk worth knowing: because the conv branch always applies ReLU before
BatchNorm1d, and the BN branch applies a second ReLU afterward (for every non-final block),
every intermediate GCN block is effectively ReLU → BN → ReLU, not the more common Conv → BN → ReLU.
BatchNorm1d itself, in training mode, standardizes per-feature over the batch (here: over all N
nodes in the full-batch graph) and applies a learned affine transform:
μ_c = (1/N) Σ_n A1^{(i)}_{n,c}
σ²_c = (1/N) Σ_n (A1^{(i)}_{n,c} - μ_c)²
B1^{(i)}_{n,c} = γ_c · (A1^{(i)}_{n,c} - μ_c) / sqrt(σ²_c + ε) + β_c
with running mean/var updated by momentum (default 0.1) for later eval-mode use, and γ, β
learnable per-channel scale/shift (also updated by the optimizer this step).
Plain dense layer, no mask effect (all-ones), no ReLU applied (final block, so the
i < len(self.layers)-1 guard is false), no BN:
logits = H^{(L-1)} (W_cls)^T + b_cls ∈ R^{N × C}
binn/train/loop.py::_loss_fn uses F.cross_entropy for classification, evaluated only on
train_mask rows:
p_{n,c} = softmax(logits_n)_c = exp(logits_{n,c}) / Σ_{c'} exp(logits_{n,c'})
loss = -(1/|train_mask|) Σ_{n ∈ train_mask} log p_{n, y_n}
Autograd differentiates through the reparametrization W^{(i)} = W^{(i)}_orig ⊙ M^{(i)}, so:
∂loss/∂W^{(i)}_orig = (∂loss/∂W^{(i)}) ⊙ M^{(i)}
i.e. any gradient signal is zeroed at every position pathway biology forbids, before it
reaches the free parameter. Gradients propagate back through the same Ŝ, BatchNorm1d
(chain rule through γ_c/sqrt(σ²_c+ε) including the batch-statistics dependency on every
sample), and ReLU (zero gradient where the pre-activation was ≤ 0) used in the forward pass —
standard backprop, just annotated here for completeness of the one-step trace.
torch.optim.Adam(model.parameters(), lr=lr, weight_decay=w_decay), defaults
lr=0.05, w_decay=5e-2, β1=0.9, β2=0.999, ε=1e-8. For every parameter θ (this includes
W^{(i)}_orig, all biases, and BN's γ, β) at step t:
g_t = ∂loss/∂θ + w_decay · θ_{t-1} — L2 term folded into the gradient (this is
plain Adam's coupled weight decay, NOT AdamW's
decoupled variant)
m_t = β1 m_{t-1} + (1-β1) g_t — first moment (mean)
v_t = β2 v_{t-1} + (1-β2) g_t² — second moment (uncentered variance)
m̂_t = m_t / (1 - β1^t) — bias correction
v̂_t = v_t / (1 - β2^t)
θ_t = θ_{t-1} - lr · m̂_t / (sqrt(v̂_t) + ε)
Consequence for pruned entries: at masked positions M^{(i)}_{p,g}=0, the loss gradient is
exactly 0 (step 5), but the weight-decay term w_decay · θ_{t-1} is not masked — Adam still
sees g_t = w_decay · θ_{t-1} ≠ 0 there, so W^{(i)}_orig at pruned positions keeps decaying
toward 0 step over step even though it never influences the forward pass (always multiplied
by 0 through M^{(i)}). Harmless, just an asymmetry worth knowing if you ever inspect
weight_orig directly instead of the effective weight.
optimizer.zero_grad() clears g_{t-1} before this step's backward pass; loss.backward()
populates ∂loss/∂θ for every leaf tensor; optimizer.step() applies the update above in place.
model.eval() — freezes BatchNorm running stats (uses accumulated running mean/var instead of
batch statistics) and disables Dropout (identity)
val_loss = cross_entropy(logits[val_mask], y[val_mask]) — no_grad, no backward
fit() tracks best_val and snapshots named_parameters() + named_buffers() (not
state_dict(), since PyG's Linear.state_dict() collapses the weight_orig/weight_mask
reparametrization back into a bare weight key that won't reload into a still-pruned module)
whenever val_loss < best_val - min_delta; patience counts consecutive non-improving epochs and
breaks the loop at patience. At the very end, the best-snapshot tensors are copied back into
the live model in place (_restore), overwriting whatever the last epoch left behind.
X (N, u_0) input features
W^{(0)}_orig, M^{(0)} (u_1, u_0) pathway-masked weight for genes→pathways
Ŝ (N, N) normalized self-looped adjacency (fixed for the epoch)
G^{(0)} = Ŝ (X W^{(0)T}) + b^{(0)} (N, u_1)
H^{(1)} = Dropout(ReLU(BN(ReLU(G^{(0)})))) (N, u_1)
W_cls, b_cls (C, u_1) dense classifier head
logits = H^{(1)} W_cls^T + b_cls (N, C)
loss = CE(logits[train_mask], y[train_mask]) scalar
One fit() epoch = one evaluation of this full pipeline forward, one loss.backward(), one
optimizer.step(), optionally one no-grad validation forward — repeated up to epochs times or
until patience triggers.