Expand description
The SAE reconstruction row as a single Taylor-jet program (issue #932).
§The row program
The exact-LAML SAE engine needs, per row, the derivative tower of the reconstruction
ẑ_row,c(p) = Σ_k ζ_k(ℓ) · decoded_{k,c}(t_k), decoded_{k,c}(t) = Σ_b Φ_b(t)·B_{b,c}— a gate nonlinearity ζ(ℓ) (softmax / IBP sigmoid) composed with a
basis Φ(t) composed with a linear decoder B, in the per-row
primary coordinates p = (gate logits ℓ, latent coordinates t). Today the
arrow-Schur assembly (SaeManifoldTerm::row_jets_for_logdet) hand-packs the
first/second channels of this reconstruction from separate gate
derivative arrays (gate_derivatives_for_row) and basis jet tensors —
exactly the kind of hand-maintained cross-block tower whose sign flips are
the #736 / desync bug genus. The #1006 third-order logdet adjoint
Γ_a = tr(H⁻¹ ∂H/∂θ_a) is the consumer of those very channels.
This module writes that reconstruction once over the
Tower4<K> scalar so the
value/gradient/Hessian/third channels of one row come from ONE jet
evaluation. SaeReconstructionRowProgram is generic over the gate kind
and the per-row basis jets; the gate, basis and decoder compose with plain
Tower4 arithmetic, so there is no separate “channel” to forget.
§The basis as a local jet
The production assembly does NOT re-evaluate the manifold basis Φ as a
function of perturbed coordinates: it consumes the precomputed jet tensors
(Φ, ∂Φ/∂t, ∂²Φ/∂t²) evaluated at the current t. The reconstruction’s
dependence on t is therefore defined by those tensors — the local
quadratic Taylor model of Φ about the current point. This program builds
each basis function as exactly that Tower4 quadratic from the stored jets,
so the value/first/second channels it emits are the same object the hand
path packs — derived by independent arithmetic (tower Leibniz / Faà di
Bruno vs hand-summed cross terms). Agreement across both is a true
correctness proof of the hand kernel; disagreement names a dropped or
sign-flipped cross block loudly. That oracle is the riding test below.
Structs§
- Atom
RowBasis Jet - One atom’s local basis jet at the current row: the stored
(value, jacobian, second)jet tensors ofΦplus the decoder blockB. Indexed[basis_col],[basis_col][axis],[basis_col][axis_a][axis_b], and[basis_col][out_col]. - SaeReconstruction
RowProgram - One row of the SAE reconstruction as a jet program: the per-atom basis jets,
the gate, the current gate-logit values, and the primary layout that maps
(atom logit, atom latent axis)to a seeded tower variable slot.
Enums§
- RowGate
- The gate nonlinearity
ζ(ℓ)of the SAE assignment, as the row program sees it. The production term carries the same two smooth branches (softmax over a shared partition; per-atom IBP/JumpReLU sigmoid); the program reproduces the branch the criterion evaluates so the value channel is the production gate.
Constants§
- SAE_
FIXED_ COORD_ SLOT - Sentinel in
SaeReconstructionRowProgram::coord_slotfor an atom coordinate that is fixed in this row’s local chart (compact active-set rows omit inactive atom coordinates, but softmax logit derivatives can still see that atom’s decoded value as a constant).