pngwn/linearizer-minimal-repro

Minimal Repro: "Who Said Neural Networks Aren't Linear?" (the Linearizer, arXiv:2510.08570)

A ~250-line 2D toy that proves the paper's core principle end to end. Not a full-paper reproduction (no LPIPS, no image diffusion, no large-scale training) — just the smallest experiment that demonstrates:

1. 1.Induced linearity is exact. f(x) = g⁻¹(A·g(x)) satisfies f(a·x+b) = a·f(x) + b·f(0) to machine precision when the identity is evaluated in the space where the algebra happens (g-space), in float64: abs err 4.4e-16, rel err 1.9e-14.
2. 2.N diffusion steps ≡ one matrix product. Composing the paper's collapsed operator B = ∏ₜ (I + Δt·(A_t − I)/(1−t)) over 100 steps matches running 100 Euler steps pointwise with MSE 2.5e-13 (1000 steps: 1.7e-12) — to float32 machine precision.
3. 3.One-step sampling works. Applying B once to g_x ~ N(0,I) and inverting through g produces samples covering all 8 modes of the target distribution with near-uniform coverage (one-step counts [580, 522, 442, 521, 573, 551, 455, 452] vs. data counts [547, 493, 481, 497, 515, 509, 563, 491] for n=4096).

Files

• —repro.py — the entire thing (training + all three checks + figure), runnable via uv run repro.py (PEP 723 deps) or plain python repro.py.
• —metrics.json — raw numbers from the run.
• —linearizer_toy.png — data vs 100-step vs one-step samples, plus the loss curve.

(The trained checkpoint.pt is not stored here — rerun repro.py to regenerate it; ~20 minutes on CPU.)

Setup

• —g: RealNVP-style invertible network (8 affine couplings, hidden 64, ActNorm, 2D data) — the small analog of the paper's invertible g.
• —A_t: rank-4 LoRA core, A_t = I + U·diag(a(t))·Vᵀ, with scalar time-dependence per rank via a small MLP on a sinusoidal time embedding — the small analog of the paper's time-dependent linear operator.
• —Data: 2D ring of 8 Gaussians (radius 2, σ=0.12), batch 256, 8000 steps at lr 1e-3.
• —Loss: the paper's induced-space flow-matching loss — all in g-space: predict the velocity g(x₀) − g(x₁) from g(x_t) with x_t = (1−t)·x₁ + t·x₀, x₁ ~ N(0,I). LPIPS replaced by MSE since there are no images here; the induced-space structure (what the paper actually linearizes) is preserved exactly.
• —Hardware/cost: ran on throttled CPU in ~20 minutes, effectively free.

Structure follows the official implementation

Loss, sampler, and collapse-matrix construction copied in structure from one_step/train_one_step.py (commit adfca2c) of the official repo (assafshocher/Linearizer); the 2D modules (InvertibleG, LinearCore) are the 2D analogs of the official modules/linear_network.py components.

What the checks do and do not show

• —The N-step ≡ 1-step collapse (check 2) is the paper's headline claim and holds exactly here, as it must: it follows algebraically once the sampler is a composition of (affine-in-g) maps. This repro confirms the implementation actually realizes that algebra, rather than the claim being an approximation in practice.
• —Induced linearity (check 1) likewise holds exactly, but note that evaluating it in data space through g⁻¹ amplifies float round-trip error because g is ill-conditioned at this scale — the identity is exact where the algebra happens.
• —Sample quality (check 3) is the only "learned" result and is accordingly the weakest: it demonstrates the principle (all modes reached in one step) at toy scale, not image quality parity with the paper.

Reference

• —Paper: arXiv:2510.08570 — Berman, Hallak, Shocher, "Who Said Neural Networks Aren't Linear?"
• —Official code: assafshocher/Linearizer