25  Lab 09 — ReLU-Sheaf Heat Equation

Anchor chapter: Chapter 9 — Neural Sheaf Heat Equation (Identity Output).

Goal. Integrate the state-dependent heat equation \(\dot{x} = -\mathcal{L}_{\mathcal{F}_t} x\), and track the times at which activation patterns switch.

Integrate the sheaf heat equation (Def. 9.1) for the [2, 4, 1] network from a random cochain initialisation. Use a simple explicit Euler scheme inside each activation region and detect region crossings by sign changes in \(c_{z^{(\ell)}}\). Plot (a) the trajectory’s projection onto two hidden-layer coordinates, (b) Dirichlet energy \(V(t)\) on a semi-log axis, and (c) a time-axis banded by the active \(\sigma\)-pattern. Confirm the exponential envelope predicted by Thm. 9.4 and locate the kinks at switching times.

TipRuns in your browser

This lab requires only NumPy, Matplotlib, and SciPy, loaded automatically via Pyodide. Code cells run directly in the page via WebAssembly — no local Python installation needed.

Prefer a local Jupyter environment? Download lab-09-relu-heat-equation.ipynb

25.1 Setup

25.2 1. Build the object

We represent the cochain as \(c = (x_{\text{in}}, z^{(1)}, z^{(2)}) \in \mathbb{R}^{n_0+n_1+n_2}\), where \(x_{\text{in}}\) is pinned (the input) and the free part is \((z^{(1)}, z^{(2)})\). The state-dependent Laplacian \(\mathcal{L}_{\mathcal{F}}(\sigma) = \delta_\Omega(\sigma)^\top \delta_\Omega(\sigma)\) depends on the activation pattern \(\sigma = \mathbf{1}[z^{(1)} > 0]\), which changes whenever any pre-activation crosses zero. A region crossing is detected by comparing \(\text{sign}(z^{(1)})\) before and after each Euler step.

25.3 2. Verify a theorem / run an experiment

We integrate \(\dot{c}_\Omega = -\mathcal{L}_{\mathcal{F}}(\sigma(c)) c_\Omega + \delta_\Omega(\sigma)^\top \tilde{b}\) with forward Euler, using a small fixed step \(\Delta t\) and updating \(\sigma\) at each step. When a region crossing is detected (a sign change in \(z^{(1)}\)), we log the time and the new pattern. Theorem 9.4 predicts that the Dirichlet energy decays as \(e^{-2\lambda_{\min} t}\) within each activation region; the kinks in the semi-log plot mark switching events.

25.4 Exercises

  1. Counting switches. Run the integration 20 times with different random initialisations and record the number of activation-pattern switches before convergence. Plot a histogram. Does the number of switches depend on the magnitude of the initial cochain?

  2. Kink sharpness. At each switching time, zoom into the Dirichlet energy plot and measure the change in slope \(d(\log V)/dt\) before and after the switch. Does the slope always decrease (energy decays faster after a switch)? Is this consistent with Theorem 9.4?

  3. Wider hidden layer. Repeat the experiment with a [2, 8, 1] network. How does the number of switches and the convergence rate change compared to [2, 4, 1]?

  4. Step-size and switching artefacts. With a very large step size (\(\Delta t \approx 0.1\)), the Euler scheme can “skip over” a region boundary. Implement a simple bisection to detect the exact crossing time and compare the trajectory to the coarse-step version.