Add _post_step hook to BaseGradientModel for proximal-gradient variants - #229
Closed
jameschapman19 wants to merge 1 commit into
Closed
Add _post_step hook to BaseGradientModel for proximal-gradient variants#229jameschapman19 wants to merge 1 commit into
jameschapman19 wants to merge 1 commit into
Conversation
Subclasses that need a proximal step after each momentum update (e.g. a sparsity penalty on the weights) currently have no way to add one without copy-pasting the whole _gradient_descent loop. _post_step is a no-op hook called on the weights right after the momentum update and before the convergence check, so a proximal-gradient variant only needs to override this one method -- momentum, batching, initial weights, and the tol check are all inherited unchanged. Two tests: the default hook is provably a no-op (identical fit to the unmodified base class), and an overridden hook actually fires every iteration (not just once at the end), verified via a call counter plus the expected effect on the final weights.
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
2 participants
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.
Summary
_post_step(weights)hook toBaseGradientModel, called on the weights right after the momentum update and before thetolconvergence check. Default implementation is the identity (no-op), so this is fully backward compatible — every existing model's fit trajectory is unchanged.CCA_EY(e.g. adding a sparsity penalty on the weights) currently has no way to add a proximal step without copy-pasting the entire_gradient_descentloop just to insert one line. With this hook, such a variant only needs to override_post_step— momentum, batching, the data-informed initial weights, and the convergence check are all inherited unchanged.CCA_EYsubclass instead.Test plan
test_post_step_default_is_identity: the unmodified base hook produces byte-identical weights to not having the hook at all.test_post_step_override_is_applied_during_training: an overridden hook (hard-thresholding) is confirmed to fire on every iteration (not just once at the end, which a naive test could miss) via a call counter, plus the expected effect on the final fitted weights.pytest tests/ -q→ 511 passed, 7 skipped (unrelated skips), no regressions.Generated by Claude Code