Add categorical (softmax) root layers to the vectorised backend#400
Merged
Conversation
LegrandNico
force-pushed
the
feat/natural-gradient-learning-kinds
branch
2 times, most recently
from
b8460e1 to
3abf70d
Compare
A frozen volatility level is never predicted or updated, so allocating its six LayerState fields only carries dead arrays through every tree_map over the state. LayerState.create(has_volatility_parent=False) now sets them to None; _match_child_vol_structure reconciles pytree structure where such a layer meets a LayerStack with a different volatility structure. Also computes the value-coupling variance in the volatile prediction as a matrix product with a per-parent vector instead of materialising a weight-matrix-sized intermediate, which matters under vmap over a batch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replaces the five weight-gradient modes with three on a single axis — which geometry the update respects: "standard" (no metric), "precision_weighted" (the backprop-parity mode), and "natural" (the exact natural gradient, whose Sherman-Morrison shrinkage reaches the conjugate posterior mean in one step, with Bernoulli and damped-multinomial-Fisher variants for binary and categorical children). "precision_ratio", "map_natural", and "pure_natural" are removed. Every kind factorises into a child-side and a parent-side vector; vectorized_weight_gradient_factors exposes the two factors so a batched caller can average gradients over samples with one contraction instead of materialising a weight-matrix-sized gradient per sample. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
LegrandNico
force-pushed
the
feat/categorical-layers
branch
from
d4b5af8 to
bccd565
Compare
LegrandNico
changed the title
LegrandNico
force-pushed
the
feat/categorical-layers
branch
from
bccd565 to
150205c
Compare
Codecov Report❌ Patch coverage is
🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
A categorical layer is one softmax choice across its nodes: the prediction normalises the coupled parent activations into class probabilities, and the prediction error is the zero-sum residual
one_hot - p. The belief-propagation sweeps route a categorical leaf through the new update functions.