Per-sample nested-Diagonal Dz/DM_Dzz (batch structure, Phases 1-3)

[a-rahimi]

Summary

Exploits the verified per-sample (batch) block-diagonal structure of the activation blocks. Sample i's layer output depends only on sample i's input, so each layer's Dz and DM_Dzz are batch independent w×w sub-blocks. Instead of materializing the dense (batch·w)² tensor and discarding the zeros, we vmap a per-sample jacobian/hessian over the batch and return a nested bpm.Diagonal.

Implements Phases 1–3 of docs/batch-structure-plan.md; the arrowhead solver (Phase 4) is deliberately not included — it is gated on the Phase 3 benchmark and a decision (see below).

Changes

• src/hessian.py: BlockWithMixedDerivatives.derivatives() computes Dz/DM_Dzz per-sample via torch.func.vmap, returned as nested bpm.Diagonals. Gated by a batch_structured_dz class attribute — True for per-sample layers, False for LossLayer (its scalar mean-loss output has no per-sample axis; its blocks stay dense).
• src/block_partitioned_matrices.py: additive Diagonal dispatch registrations so nested Diagonals flow through K assembly and both solvers unchanged; fixes a latent .shape[1]→.width bug in UpperDiagonal @ Diagonal.
• Both solver="splu" and solver="block" are untouched and remain as references.

Results (independently verified)

• Correctness: structured output matches a float64 dense (H+εI)⁻¹g reference to 2e-6 (splu) / 1e-3 (block).
• Speed/memory (batch 32, SequenceOfDenseBlocks(768, 8, 10, num_layers=8)): ~2× faster setup and ~2× less peak RSS. The default batch 128 — previously OOM-killed by the ~38 GB dense DM_Dzz — now completes. Full table in docs/phase3-benchmark.md.
• Tests: 136 pass (133 + 3 new for the nested-Diagonal representation).

Phase 4 go/no-go context

The Phase 3 benchmark surfaces a strategic point: the win here is entirely in setup (the derivatives pass); K assembly and the splu solve are unchanged. The arrowhead solver (Phase 4) would speed up the solve, which is currently the minority of the budget at these sizes — so its value proposition is weaker than expected. Recommend treating Phases 1–3 as a standalone win and deciding Phase 4 separately.

Coordination

Branched off main, not off PR #3 (splu-timing). PR #3 will need rebasing on top of this once merged — both edit hessian.py around the derivatives pass / K assembly (see the plan's coordination note).

🤖 Generated with Claude Code

[arahimi-hims]

don't need this file as part of the PR. it's ephemeral

[arahimi-hims]

as long as these benchmark results are in the pr description, don't need this file

[arahimi-hims]

are you doing this so that structured matrices can multiply unstructored tensors? if so, we don't want this. we want the rhs to have a structure that matches the lhs.

[a-rahimi]

Addressed all three comments (force-pushed a cleaned single commit):

• Dropped docs/batch-structure-plan.md and docs/phase3-benchmark.md — ephemeral; benchmark results stay in the PR description.
• Removed the Diagonal-times-dense-Tensor chunking (both __matmul__ dispatches, the Diagonal.solve(Tensor) branch, and the _split_dense_by_sizes/_matmul_dense_rhs_blockwise/_solve_dense_rhs_blockwise helpers). No structured matrix multiplies an unstructured tensor anywhere now — Diagonal @ Tensor raises.

On "we want the rhs to have a structure that matches the lhs": in the block and matrix-vector reference paths the per-sample structure genuinely can't survive the solve, so the fix is to densify those activation blocks rather than structure the RHS. Verified numerically — the block solver eliminates the shared parameters first, and the resulting param-border Schur term A10·A00⁻¹·A01 in activation space is purely cross-sample (within-sample mass 0.000, cross-sample 83.7 on a small case): elimination through the shared-parameter border couples every sample, so the matrices actually factorized (S11, S22) are dense across samples. A structure-preserving solve requires eliminating activations first — i.e. the arrowhead solver (deferred Phase 4).

The default splu path keeps the full per-sample structure via to_scipy_csc, so setup's ~2× memory/time win is unaffected. Both solvers still match a float64 dense reference (splu 2e-6, block 1e-3); 136 tests pass.