feat: add projection_dim support to kronecker factored Hessians - #337
Merged
Conversation
luciaquirke
force-pushed
the
feat/compressed-kfac-apply
branch
from
July 18, 2026 17:03
edb08d1 to
59fe5e1
Compare
EKFAC's H^-1 must be applied to the full, unprojected gradient (the
eigenbasis rotation only makes sense at the true parameter dimension), so
apply_hessian.py's IVHP path and the query build step have always forced
projection_dim=0. But the training-side gradient store built by `bergson
build`/`bergson score` uses index_cfg.projection_dim's per-module Kronecker
random projection by default, so the hessian_pipeline's scoring step had to
force projection_dim=0 there too -- silently ignoring whatever compression
the user asked for and scoring against full dense gradients.
EkfacConfig gains projection_dim/projection_type; when set,
EkfacApplicator.compute_ivhp_sharded compresses each module's H^-1 G output
to [p, p] as a post-processing step, using the same create_projection_matrix
identifiers (f"{name}/left" / f"{name}/right") that bergson build already
uses on training gradients. The hessian_pipeline now threads
index_cfg.projection_dim/projection_type into the apply step and no longer
overrides the score step's projection_dim to 0, so a compressed query and a
compressed training-side store are directly comparable end to end.
Supersedes #275, which predates the FactoredPreconditioner
refactor and no longer applies; ports the same idea (compress the IVHP
output rather than the pre-Hessian gradient) onto the current apply path.
Co-Authored-By: Girish Gupta <girish@girishgupta.com>
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
luciaquirke
force-pushed
the
feat/compressed-kfac-apply
branch
from
July 18, 2026 17:06
59fe5e1 to
157fa1b
Compare
…ssian combos - EkfacApplicator now rejects projection_dim>0 combined with ev_correction=True (EK-FAC): the post-projection compression path hasn't been validated for the corrected, per-cell eigenvalue grid EK-FAC produces, only plain K-FAC's outer product. - hessian_pipeline now rejects preprocess_cfg.unit_normalize (cosine similarity): the pipeline only ever fits and applies a Kronecker-factored Hessian, and split/cosine scoring for factored Hessians is only supported through the existing bergson score --hessian_path path (uncompressed). Co-Authored-By: Girish Gupta <girish@girishgupta.com> Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…t_normalize score_dataset already rejected any Kronecker-factored hessian combined with index_cfg.projection_dim != 0, regardless of unit_normalize -- necessary since the query itself gets H^-1 applied unconditionally and needs full, unprojected dims. Give the unit_normalize=True case (cosine similarity / split H^-1/2 scoring, which additionally routes the index side through preconditioner.apply) its own message naming that as the cause, and pin the combination with a regression test. Co-Authored-By: Girish Gupta <girish@girishgupta.com> Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
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.
Based on #275, rebased and updated on the recent code changes
This is H^-1 support: H^-{1/2} will require computing the hessians in the projected space which is uncharted territory for the kronecker factored Hs
cc: @jammastergirish