Integrate LieBNSPD into spd_learn#21
Merged
bruAristimunha merged 22 commits intospdlearn:mainfrom Mar 22, 2026
Merged
Conversation
Integrate LieBNSPD as a package module and update the LieBN examples and tests to import the shared implementation from spd_learn.modules. This commit is based on Bruno Aristimunha's initial integration draft. I reviewed it and adjusted the package wiring while adding and refining comments/docstrings around the LieBN implementation and its provenance.
Extend karcher_mean_iteration with an optional tangent return and use it in LieBNSPD to stop AIM Karcher iterations early once the tangent mean norm is sufficiently small.
Set test_liebn to use float64 by default and increase the simulated SPD dimension to keep the LieBN test configuration stable under double-precision runs.
- Rename LieBNSPD → SPDBatchNormLie to match naming convention (SPDBatchNormMean, SPDBatchNormMeanVar) - Add device/dtype params to SPDBatchNormLie.__init__ - Use buffer rebinding instead of .copy_() for running stats, matching the pattern of other batchnorm layers and enabling broadcast over arbitrary leading batch dimensions - Fix Karcher early stopping for multi-batch tensors (.max()) - Add SPDBatchNormLie to public API exports and docs autosummary - Add mandatory params and complex dtype skip in integration tests - Fix __file__ NameError in sphinx-gallery example - Add geoopt to brain dependencies for TSMNet example - Add chen2024liebn bibtex entry to references.bib - Add rtol=0.05 to variance convergence test for LEM/LCM
Extract reusable math from the module into functional/batchnorm.py: - spd_cholesky_congruence: Cholesky-based congruence action (L X L^T) - lie_group_variance: weighted Fréchet variance for AIM/LEM/LCM The module now delegates to existing functional operations: - cholesky_log/cholesky_exp for LCM deform/inv_deform - log_euclidean_scalar_multiply for AIM variance scaling - matrix_log/matrix_exp/matrix_power for AIM/LEM deform - karcher_mean_iteration for AIM Fréchet mean
Support two implementations of the AIM group action: - "cholesky": L⁻¹ X L⁻ᵀ via Cholesky factor (original LieBN paper) - "eig": M⁻¹/² X M⁻¹/² via eigendecomposition (spd_centering) Both are valid SPD batch normalizations that center the mean to identity, but use different geometric transports. Default remains "cholesky" to match the paper.
The TSMNet tutorial depends on geoopt (RiemannianAdam) which is not a project dependency. Remove the tutorial and the geoopt entry from the brain extras.
P1: Validate metric parameter in SPDBatchNormLie.__init__
P1: Move .detach() from lie_group_variance to caller (keep
functional API stateless)
P1: Move ensure_sym import to top-level in batchnorm.py
P2: Document Karcher convergence threshold (1e-5) in docstring
P2: Show congruence in extra_repr
P2: Validate metric in lie_group_variance
P2: Fix torch.set_default_dtype leak in test_liebn.py (autouse fixture)
P3: Alphabetize SPDBatchNormLie in __all__
P3: Group SPDBatchNorm* together in api.rst
P3: List new functions in batchnorm.py module docstring
Also add congruence parametrization to test_post_normalization_mean
for coverage of both cholesky and eig paths.
Replace the autouse fixture that set global default dtype with explicit dtype=torch.float64 passed to every SPDBatchNormLie constructor and to data generation tensors. This avoids leaking global state to other test modules and properly exercises the dtype parameter.
Co-authored-by: Bru <b.aristimunha@gmail.com>
Add frechet_mean() to spd_learn.functional.batchnorm, unifying the duplicated Karcher flow logic from SPDBatchNormMean, SPDBatchNormMeanVar, SPDBatchNormLie, and the SPDIM tutorial into a single reusable function. Rename LieBN.py to liebn.py for snake_case consistency with all other module files, and rename karcher_steps to n_iter in SPDBatchNormLie to match the other batchnorm modules.
…tency All other batchnorm modules (SPDBatchNormMean, SPDBatchNormMeanVar, BatchReNorm) use `num_features` as their matrix-size parameter. This aligns SPDBatchNormLie with the same convention.
Contributor
Docs CI build analysisThe docs build is hanging for 6+ hours (vs ~23 min on main) because Issues found1. Training parameters are too heavy for CI The example runs n_clusters = 3 # Reduced from 5 for faster documentation build
max_epochs = 20 # Reduced from 100 for faster documentation buildSuggested fix — detect CI and reduce: import os
EPOCHS = 5 if os.environ.get("CI") else 200
N_RUNS = 2 if os.environ.get("CI") else 102. Line 291: CHECKPOINT_PATH = Path(tempfile.gettempdir()) / "liebn_checkpoint.json"3. The failed build also showed CI improvements on
|
2 tasks
…dtype(torch.float64) Several examples set torch.set_default_dtype(torch.float64) globally, which persists across sphinx-gallery examples in the same worker process. This caused MATT and GREEN examples to fail with dtype mismatches when their models were initialized with float64 parameters but braindecode cast inputs to float32. Changes: - Remove torch.set_default_dtype(torch.float64) from 4 example scripts - Add explicit .float() conversion where numpy float64 data enters torch - Add _reset_torch_defaults to sphinx-gallery reset_modules as safety net - Add howto subsection to sphinx-gallery ordering
…amples The docs CI was taking ~6 hours because every push re-executed all gallery examples (data downloads + model training). Sphinx-gallery already skips unchanged examples locally via MD5 checks on the generated/ directory, but CI never persisted this cache. The cache key hashes all example source files so only modified examples re-execute on subsequent runs.
…EADME Add a new "Batch Normalization on SPD Manifolds" section to geometric_concepts.rst covering why Euclidean BN fails, Riemannian BN, and the LieBN 5-step pipeline with metric comparison table. Add README.txt for the new howto gallery section.
…itZH-Chen-main
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
This PR integrates LieBNSPD into
spd_learn.modulesas a shared package moduleand updates the related examples and tests.
It also refines the AIM Karcher mean computation in
LieBNSPDby adding anearly stopping condition (which can be used in other AIM-based Karcher calculation), and updates the LieBN
tests to run in double precision.
Main changes
spd_learn/modules/LieBN.pyas the sharedLieBNSPDimplementationLieBNSPDfromspd_learn.modulestests/test_liebn.pyto import theshared package implementation
LieBNSPDkarcher_mean_iterationwith optional tangent return support forconvergence checking
tests/test_liebn.pyto usefloat64and increase the simulated SPDdimension for more stable double-precision testing