Precision weighting + NGD trainer for predictive coding (ResNet-18/CIFAR-10)

[aprotsenko24]

Precision weighting + natural-gradient (NGD) for predictive coding

Adds uncertainty-aware (precision-weighted) predictive coding to FabricPC and benchmarks it on the
ResNet-18 / CIFAR-10 muPC demo. Precision ≈ inverse error covariance ≈ Fisher information, so this is
a mechanism-grounded conditioning lever (complementary to muPC scaling).

What's added

• GaussianEnergy: per-channel diagonal precision (scalar path unchanged) + optional report-only
  log-precision term.
• NodeState.precision: optional per-node precision carried in state (default None ⇒
  bitwise-identical to prior behaviour, no extra pytree leaves). Consumed at the single
  energy_functional choke point, so a dynamic precision reaches both inference and the local
  weight gradients with no node forward() / inference changes. muPC weight_grad_scale=1.0 ⇒ no
  double application.
• core/precision.py::probe_residual_precision — one-pass diagonal Π = 1/Var(residual).
• training/ngd_trainer.py — train_ngd (online EMA precision; NGD via precision-weighted SGD),
  probe_latent_propagation (per-layer signal-propagation diagnostic).
• resnet18_cifar10_demo.py — --precision {none,diag_probe,online}, --optimizer {adamw,ngd},
  --momentum, --seed, --probe_latents; machine-readable [RESULT] line.
• precision_experiment.py — paired multi-seed sweep driver (CSV checkpoint, --resume, per-epoch
  progress capture, retries, optimizer/momentum-keyed paired-stats report).
• Tests: tests/test_precision_weighting.py, tests/test_ngd_precision.py.

Results (CIFAR-10, muPC ResNet-18) — two separate modes

Mode A — diag_probe (frozen diagonal precision, AdamW). Π=1/Var estimated once pre-training,
then fixed.

• Consistently better in EARLY training (2 epochs, every seed): +1.57 / +2.84 / +3.14 pts over
  the Π=1 baseline, with lower train energy in all three.

• But the gain washes out at convergence. 30 epochs, 3 paired seeds:

  seed	none	diag_probe	Δacc
  0	42.17%	40.71%	−1.46
  1	42.16%	42.87%	+0.71
  2	41.77%	42.45%	+0.68

  none 42.03 ± 0.23% vs diag_probe 42.01 ± 1.15% → paired Δ = −0.02 pts, p = 0.977 (not
  significant); train energy marginally lower for diag_probe (2/3 seeds).

• Interpretation: precision weighting is a convergence-speed / conditioning effect — it
  accelerates early training but does not change the converged accuracy. (Likely the frozen Π,
  estimated from the untrained net, goes stale as the error structure evolves.)

Mode B — online / NGD (per-step EMA precision, plain SGD). Π updated every step; SGD so precision
is the per-channel adaptive LR. Finding: SGD/NGD trains far slower than AdamW here (narrow stable-LR
window; momentum helps but doesn't match Adam in equal epochs) and does not reach the AdamW baseline —
reported as a negative result.

Honest framing: precision weighting is a conditioning/robustness/convergence-speed lever. It is
not a path to backprop-level accuracy — PC-ResNet-18 plateaus ~42% (backprop ResNet-18 ≈90%); that
ceiling is a property of PC training on this task, which precision does not move.

Tests

Full suite 229 passed (the NodeState pytree change is a clean no-op regression); the 2 new files
add 15 tests (diagonal/scalar energy, NodeState.precision backward-compat, precision reaching the
weight gradients, online EMA update, normalize-after-clip invariant, latent probe). Run with the
venv (CUDA→CPU fallback is normal): PYTHONPATH=. python -m pytest tests/ -q.

Reproduce

PYTHONPATH=. python examples/precision_experiment.py --seeds 0 1 2 --arms none diag_probe \
    --num_epochs 30 --activation tanh --no_augment --eval_every 5 --resume \
    --out precision_30ep.csv --md precision_30ep.md

Follow-ups (not in this PR)

• PC-vs-backprop reference run for the ~upper bound.
• 30-epoch online/NGD (Mode B) paired sweep at a tuned SGD LR.