Repository navigation
Audit + full verification: 0.4.0 (scalar second moment, external oracles) - #2
Merged
Merged
Conversation
A resume=False reinitialization over an existing directory rewrote w/m/v/t directly, without the reinit_in_progress marker. A crash in the middle of _write_fresh_state (e.g. kill -9 after w was written) left a mixed state - new w, old m/v/t, still-valid old meta - that the next resume=True accepted silently: all state files have the same size, so the size verification cannot detect it. 0.2.0 had the marker; the 0.3.0 constructor path dropped it (import_checkpoint kept it). The marker is now written before the data files are touched whenever the reinitialization runs over an existing directory; resume=True refuses a marked state with the recovery instructions, resume=False performs the documented full recovery and clears the marker. The crash is simulated in the test by failing _write_fresh_state right after the w write.
A negative per-row counter on disk (torn write, corrupted t.dat) survived row validation: the fresh-row branch zeroed m/v but kept the counter, so bias correction became 1 - beta**t with t < 0 - negative - and sqrt(v/bc2) produced NaN, aborting the whole step with a misleading 'non-finite weights; reduce lr' error instead of the documented A4 row reset (verified by a counter-proof run before the fix). The corruption check now includes t_prev >= 0: such a row is fully reset (w=m=v=0, t=0) and the step proceeds as a canonical first Adam step from the origin (pinned analytically by the new test).
weight_decay_mode is stored in per-parameter meta and in checkpoint manifests, but a mismatch with the constructor argument was never reported: a state trained with the coupled "l2" penalty could silently continue with "decoupled" (AdamW-style) decay, or vice versa, changing the meaning of every subsequent step. resume=True and import_checkpoint now log a warning naming both modes when they differ.
The module docstrings claimed the hyperbolic step 'implements the Riemannian Adam of Becigneul & Ganea (ICLR 2019)' and the changelog called the sqrt-transport of v 'the paper's operation'. Checked against Algorithm 1 of the paper and geoopt's RiemannianAdam: the paper keeps the second moment scalar (the squared Riemannian norm, which needs no transport) and transports only the first moment (the tau update); geoopt keeps per-coordinate moments and does not transport v at all. DSA uses per-coordinate moments (geoopt style), transports m as in the paper, and transports sqrt(v) squared back - a positivity-preserving heuristic that keeps the Riemannian scale of sqrt(v). Docstrings, the reference module and the 0.3.0 known-limitations note now say exactly this; the writer also documents that per-row write ordering is guaranteed through step() only. Version 0.3.1; CHANGELOG entry lists the three audit fixes with their regression tests.
The hyperbolic second moment is now the scalar squared Riemannian norm lambda(x)^2*||g_riem||^2 per row (Algorithm 1 of Becigneul & Ganea; geoopt's component_inner - the ball is one manifold component), stored broadcast over coordinates. The 0.3.x per-coordinate accumulator was not intrinsic (chart-rotation dependent) and, without the metric weight, produced first steps whose Riemannian length grows as lambda(x)*lr near the boundary instead of the bounded lr*sqrt(dim) - pinned by the new invariant test. A scalar needs no parallel transport, so the sqrt-transport heuristic is removed: only the first moment is transported (the paper's tau update). The Euclidean k=0 path keeps per-coordinate moments (R^d is a product of d components). State format v3 / checkpoint v3: v1/v2 moments are incompatible and are reset on migration (weights preserved, legacy_migrated_from recorded); 0.3.0 (v2) directories migrate with the same rule - covered by a new test.
test_geoopt_step.py compares the FULL multi-step trajectory (w, m and v
over 10 steps; k in {0.7, 1.0}; weight decay in {0, 0.02}) against
geoopt 0.5.1's RiemannianAdam on PoincareBallExact - agreement to
float32 rounding (~2e-7). Runs when geoopt is importable (install into
a separate dir + PYTHONPATH), skips cleanly otherwise, e.g. in CI.
test_torch_equivalence.py pins the Euclidean k=0 path against
torch.optim.Adam (l2 mode) and torch.optim.AdamW (decoupled mode) over
dense 12-step trajectories within 2e-6 - an oracle independent of this
repository.
test_optimization.py: 300 steps minimise the geodesic distance to fixed targets on the ball (from ~1.5 to <0.05 per row, autograd gradients through the verified distance formula); a near-boundary stress test (rows at 0.9999 of the radius, gradient scale 50) stays finite and inside the ball. README step description updated for the scalar second moment; CHANGELOG 0.4.0 entry records the mathematics change and every oracle; version bumped in pyproject, dsa.__init__ and the packaging test.
The package job hard-coded 'Version: 0.3.0' in the wheel-metadata check, so every version bump broke CI (0.4.0 did). The check now reads the version from pyproject.toml — single source of truth, same guarantee.
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
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.
Line-by-line audit of 0.3.0 against the 0.2.0 baseline, Algorithm 1 of Becigneul & Ganea and geoopt, followed by full verification against independent oracles.
Fixed (audit findings)
Changed (mathematics)
Verified against independent oracles
Suite: 144 passed locally with geoopt, 140 passed + 1 skipped without (CI).