28  Lab 12 — Sheaf Trainer

Anchor chapter: Chapter 12 — Backprop-Free Training, Batches, and Timescale Separation.

Goal. Implement the joint (cochain, parameter) gradient flow on the paraboloid \(f(x,y) = x^2 + y^2\) and compare convergence with vanilla SGD.

Implement Algorithm 12.4 end-to-end on the [2, 30, 1] paraboloid task. Train with sheaf-based updates and an SGD baseline, both with the same initialisation, batch size, and learning rate. Plot training loss vs epoch for both, measure wall-clock time, and confirm the \(1/\lambda_{\min}^{\text{free}}\) slowdown predicted by Thm. 12.5. Extend the lab with a batch-parallelism demonstration: run \(B\) fast-phase cochain solves concurrently (e.g., via joblib) and time the speedup.

TipRuns in your browser

This lab uses NumPy and Matplotlib for the sheaf trainer, and MyGrad — a pure-Python, NumPy-based autograd library — for the SGD baseline, in place of PyTorch (which has no Pyodide wheel). MyGrad is installed on the fly via micropip in the Setup cell below; its API (mg.Tensor, .backward(), .grad) mirrors PyTorch’s closely enough that the training loop reads the same. Everything runs directly in the page via WebAssembly, no local install needed.

Prefer a local Jupyter environment with real PyTorch? Download lab-12-sheaf-trainer.ipynb

Install dependencies: pip install torch numpy matplotlib joblib

28.1 Setup

28.2 1. Build the object

Algorithm 12.4 separates training into two timescales: a fast phase (cochain solve — given current weights, find the cochain \(c^*\) that minimises \(\|\delta_\sigma c - \tilde{b}\|^2\) for each training sample) and a slow phase (parameter update — use the solved cochains to compute the gradient of the residual energy \(R(\theta)\) and take one step in weight space). The fast phase is a single triangular back-substitution per sample; the slow phase is a standard gradient step. We implement both and compare to vanilla SGD on the paraboloid regression task \(f(x_1, x_2) = x_1^2 + x_2^2 - 2/3\).

28.3 2. Verify a theorem / run an experiment

We train both the sheaf-based algorithm (gradient on \(R\) via the fast-phase solve) and vanilla SGD (backprop on MSE) for 200 epochs with matched learning rates and report the training loss and wall-clock time. Theorem 12.5 predicts a \(1/\lambda_{\min}(L_{\text{free}})\) slowdown for the sheaf-based method relative to SGD in the initial transient; we estimate \(\lambda_{\min}\) from the assembled \(L_{\text{free}}\) at the initial weights.

28.4 Exercises

  1. Layer-wise updates. The implementation above only updates \(W_2\) and \(b_2\) (the output layer). Extend the slow phase to also update \(W_1\) and \(b_1\) by differentiating \(R\) with respect to those parameters (requires chain rule through the fast-phase solve). Compare convergence of the full sheaf trainer to the partial one.

  2. Batch parallelism. Use concurrent.futures.ThreadPoolExecutor (or joblib.Parallel) to run the fast-phase cochain solves for all \(N\) samples concurrently. Measure wall-clock time for batch sizes \(B \in \{1, 10, 50, 200\}\) and plot the speedup.

  3. Timescale separation. Theorem 12.5 predicts a \(1/\lambda_{\min}\) slowdown. Plot the ratio of initial loss-decrease rates (SGD vs sheaf) against \(1/\lambda_{\min}\) computed at the initial weights for networks with different hidden widths \(n_1 \in \{10, 30, 100\}\).

  4. Regularisation as sheaf augmentation. Adding \(\ell_2\) regularisation \(\tfrac{\lambda}{2}\|\theta\|^2\) to the loss is equivalent to adding a “reluctance edge” from each weight vertex to a zero anchor. Implement this by modifying the fast-phase RHS and verify that the regularised sheaf trainer converges to a smaller weight norm.