[michakraus]

Summary

SymplecticAttentionQ and SymplecticAttentionP with symmetric = true computed the shear increment 2 A X σ(XᵀAX) for either activation. That is the gradient of the attention potential only when σ(XᵀAX) is symmetric:

• MatrixSoftmax (the default) normalises the whole matrix, so σ is symmetric and the layer was correct.
• VectorSoftmax normalises each column. The potential Σ_j log Σ_i exp(c_ij) then has the gradient A X σᵀ + Aᵀ X σ, which is A X (σ + σᵀ) for symmetric A. With 2 A X σ the layer was not a gradient shear and not symplectic: JᵀΩJ − Ω is of order one at the default initialisation.

This PR computes A X (σ + σᵀ) for both activations:

• src/layers/symplectic_attention.jl: the :symmetric method. The single A = ps.A binding stays, so the parameter gradient is still a SymmetricMatrix.
• src/layers/inverse_layers.jl: InverseSymplecticAttention uses the same formula, so it still inverts the layer.
• docs/src/layers/symplectic_attention.md: the vector-softmax simplification for symmetric A is corrected.

The symmetric = false method already computed A X σᵀ + Aᵀ X σ and is unchanged. With MatrixSoftmax the output is unchanged up to round-off; with VectorSoftmax and symmetric = true, outputs change for the same parameters.

Tests

• test/layers/symplectic_attention.jl: a new testset checks JᵀΩJ ≈ Ω for the Zygote Jacobian over a sequence, for Q and P layers, both activations and both values of symmetric. It fails on main for VectorSoftmax, symmetric = true and passes on this branch; the other cases are controls that pass on both.
• test/layers/inverse_layers.jl: the exact-inversion test now covers the VectorSoftmax symmetric layers.
• All test files that the diff reaches pass locally (6085 tests).

Pre-PR verification

Two verifier runs: a full seven-point pass on b869c8e and 525ab39, and a scoped re-check of 8dbb8b0.

Fixed

• CHANGELOG: the old defect size was understated ("up to about 2e-2"); it measured 0.126 (Q) and 1.41 (P) with the test's seeds. The entry now says JᵀΩJ − Ω is of order one (8dbb8b0).
• Test label (nit): the @safetestset label now names both testsets of the file (8dbb8b0).

Unresolved: none.

Pre-existing, not changed here

• nit — docs/src/layers/symplectic_attention.md: the vector-softmax potential is written with the sum inside the exponential (\exp(\sum_m C_{mn})), while the derivative needs \sum_m \exp(C_{mn}).
• nit — the page derives a "one softmax" (1 + in the denominator), while the code uses the standard softmax.
• nit — the same page says "staring point" twice.

Checked and clean

• The new increment equals an independent Zygote gradient of the attention potential to ≤ 1.3e-15, for both activations, Q and P, matrix and 3-tensor inputs.
• MatrixSoftmax outputs are unchanged to ≤ 4.4e-16, so the pinned outputs still hold.
• The inverse layers recover the input to ≤ 3.3e-16.
• The parameter gradient stays a SymmetricMatrix.
• A sweep of 20 seeds per cell over d ∈ {1,…,4}, L ∈ {2,4,8}, Q and P: the old increment gives a median max|JᵀΩJ − Ω| of 0.18–2.2; the fixed layer gives ≤ 1.2e-14.
• Aqua piracy is clean with the repository's treat_as_own = [Batch]; all forward and inverse calls infer to concrete types on Julia 1.13; fatou lint is clean; test-layout.jl --check passes.

Not checked

• The GPU/JLArray path of the changed layer: no device test covers it.
• makedocs (the docs change adds no @ref and no docstring) and inference on Julia 1.12.

Release: patch. A bug fix with no API change.

🤖 Generated with Claude Code

[coderabbitai]

Warning

Review limit reached

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Next included review available in 6 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used the included review currently available.

Learn how review limits work.

Review configuration:

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: bc74b529-8878-4204-bfaa-e877fc167e63

Commits

Reviewing files that changed from the base of the PR and between 23a0228 and e0f7328.

Files selected for processing: 1

• docs/src/layers/symplectic_attention.md

No actionable comments were generated in the recent review.

Summary

Summary by CodeRabbit

• Bug Fixes
  • Corrected symmetric symplectic-attention behavior to preserve symplecticity with matrix and vector softmax activations, including inverse attention layers.
  • Matrix-softmax outputs remain unchanged except for possible round-off differences. Outputs from trained models using vector softmax may change without parameter updates.
  • Nonsymmetric attention behavior is unchanged.
• Documentation
  • Updated the symplectic-attention documentation with the corrected formula.
• Tests
  • Expanded coverage for symmetric and nonsymmetric attention configurations, both softmax activations, and layer inversion.

The symmetric symplectic-attention displacement now adds the activation to its transpose in forward and inverse layers. Documentation describes the formula, and tests check symplecticity and layerwise inversion across attention variants.

Changes: Symmetric attention formula

Layer / File(s)	Summary
Update the symmetric displacement formula src/layers/symplectic_attention.jl, src/layers/inverse_layers.jl, docs/src/layers/symplectic_attention.md, CHANGELOG.md	The symmetric forward and inverse layers add the activation to its transpose instead of multiplying the activation by two. The documentation and changelog describe the formula and its stated effects.
Validate symplecticity and inversion test/layers/symplectic_attention.jl, test/layers/inverse_layers.jl, test/runtests.jl	Tests check the canonical symplectic form for both attention variants, activation types, and symmetry settings. Inversion coverage adds nonsymmetric P attention and Q/P attention with VectorSoftmax(). The test-set label is updated.

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Possibly related PRs

• JuliaGNI/GeometricMachineLearning.jl#363: Adds inverse symplectic attention; this change updates the symmetric inverse displacement to match the forward formula.
• JuliaGNI/GeometricMachineLearning.jl#351: Consolidates symplectic-attention execution methods; this change updates the symmetric method’s increment.

Merge Risk: Minimal · previous review 23a02

The corrected attention formula is consistent across forward and inverse layers, with regression checks included. No concrete merge-blocking issue remains; the older softmax documentation inconsistency can be addressed separately.

Pre-merge checks: 5

Check name	Status	Explanation
Title check		The title clearly and concisely describes the main change: making symmetric SymplecticAttention symplectic with VectorSoftmax.
Description check		The description directly explains the formula correction, affected layers, documentation updates, and test coverage.
Docstring Coverage		No eligible functions found in the changed files. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

No eligible functions found in the changed files. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 0 functions across 0 files. (2 skipped: 2 unsupported.)

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn

CodeRabbit Verify

• Run CodeRabbit Verify to test the changes in this pull request.

Recent review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 3338f187-b685-4923-b8a6-fef16e0168e1

Commits

Reviewing files that changed from the base of the PR and between 8dbb8b0 and 23a0228.

Files selected for processing: 2

• CHANGELOG.md
• test/layers/inverse_layers.jl

Files skipped from review as they are similar to previous changes: 1

• CHANGELOG.md

Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.49%. Comparing base (1c49b5d) to head (e0f7328).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #378   +/-   ##
=======================================
  Coverage   94.48%   94.49%           
=======================================
  Files          63       63           
  Lines        2738     2741    +3     
=======================================
+ Hits         2587     2590    +3     
  Misses        151      151           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[michakraus]

JuliaGNI/GeometricMachineLearning.jl PR #378 — Make symmetric SymplecticAttention symplectic with VectorSoftmax

Verdict: approve
Reviewed: 826fe5b
CI: green on all 13 required and release-matrix checks (Julia 1 and Julia min on ubuntu-latest, macOS-latest, and windows-latest, Downgrade, Doctests, Documentation, Metal on Julia 1 and Julia min, Scripts, codecov/patch at 100%, codecov/project at 94.49%); the 2 advisory continue-on-error prerelease jobs (Julia pre and Julia nightly on ubuntu-latest) fail in julia-buildpkg while precompiling LLVM (LoadError: UndefVarError: libLLVMExtra not defined in LLVM.API).

Findings

#	severity	file:line	claim	evidence
1	nit	docs/src/layers/symplectic_attention.md:24,26,52,54,60	Pre-existing documentation defects on the same page (also flagged in coderabbitai[bot] review comment #discussion_r4454495583 and listed in the PR body under "Pre-existing, not changed here") are neither fixed in this change nor recorded in KNOWN_ISSUES.md: lines 26 and 54 place the sum inside the exponential (\exp(\sum C) instead of \sum \exp(C)), line 60 pulls 1 / (1 + \sum_m \exp(C_{mj})) outside \sum_n in the first term instead of keeping 1 / (1 + \sum_m \exp(C_{mn})) inside \sum_n, and lines 24 and 52 write "The staring point". Note on #discussion_r4454495583: the change on line 78 (\nabla_Z\Sigma(Z) = AZ(\mathrm{softmax}_1(C) + \mathrm{softmax}_1(C)^T)) is algebraically consistent with lines 66 and 72, and both standard column-wise \mathrm{softmax} (for \Sigma(Z) = \sum_n \log \sum_m \exp(C_{mn})) and \mathrm{softmax}_1 (for \Sigma(Z) = \sum_n \log(1 + \sum_m \exp(C_{mn}))) have the same gradient form A Z \sigma(C)^T + A^T Z \sigma(C) = A Z (\sigma(C) + \sigma(C)^T) when A^T = A.	Read docs/src/layers/symplectic_attention.md:24-86, src/activations/softmax.jl:27-38, and KNOWN_ISSUES.md.
2	nit	test/layers/inverse_layers.jl:87-92	@testset "exact layerwise inversion ($T)" tests (symmetric = true, MatrixSoftmax()), (symmetric = false, MatrixSoftmax()), and (symmetric = true, VectorSoftmax()), leaving out (symmetric = false, VectorSoftmax()).	Read test/layers/inverse_layers.jl:84-92; verified in a cold Julia 1.13.3 process that SymplecticAttentionQ(6; symmetric = false, activation = VectorSoftmax()) and SymplecticAttentionP(6; symmetric = false, activation = VectorSoftmax()) also invert within 10 * eps(T) for Float32 and Float64 on 1D, 2D, and 3D arrays and (q, p) NamedTuples.

Verified and fine

• Symplecticity defect on main and fix on branch: Re-measured maximum(abs, J' * Ω * J - Ω) in a cold Julia 1.13.3 process. With the test sequence in test/layers/symplectic_attention.jl (Random.seed!(1234), d = 2, L = 4), the main formula 2 * σXAX gives 0.972 (Q) and 0.102 (P), whereas the branch formula σXAX + _custom_transpose(σXAX) gives 2.22e-16 (Q) and 2.78e-17 (P). Across a sweep of d ∈ (1, 2, 3, 4), L ∈ (2, 4, 8), 5 random draws each (Q and P, 120 cases), the main formula gives cell medians in [0.309, 1.769] (overall range [2.27e-3, 5.045], median 0.932), while the branch formula gives a maximum residual of 6.66e-15 (median 3.47e-16).
• Direct potential gradient identity (2D and 3D): Verified in a cold process that _mul(_mul(ps.A, x), σXAX + _custom_transpose(σXAX)) matches Zygote.gradient(X -> sum(log.(sum(exp.(X' * Matrix(ps.A) * X); dims = 1))), x)[1] (and the batch sum for 3D inputs) with maximum absolute difference 4.44e-16 (6×4), 8.88e-16 (6×5), 3.33e-16 (6×4×1), and 1.78e-15 (6×5×3).
• MatrixSoftmax invariance: Because C = _mul(_mul(_custom_transpose(x), A), x) is symmetric per batch slice when A is a SymmetricMatrix and MatrixSoftmax normalises over (1, 2) by a scalar per slice, σXAX + _custom_transpose(σXAX) equals 2 * σXAX up to floating-point rounding (max|new - old| = 1.19e-7 = 1 * eps(Float32) for Q and 5.96e-8 = 0.5 * eps(Float32) for P in Float32; 2.78e-17 for Q and 1.11e-16 = 0.5 * eps(Float64) for P in Float64).
• Exact layerwise inversion (src/layers/inverse_layers.jl:128-134): Verified that _attention_displacement on InverseSymplecticAttention{M, M, LayerType, :symmetric} uses the identical increment _mul(_mul(A, x), σXAX + _custom_transpose(σXAX)) and satisfies l_inv(l(z, ps), ps) ≈ z and l(l_inv(z, ps), ps) ≈ z within 10 * eps(T) for all 4 (symmetric, activation) combinations, Float32 and Float64, 1D/2D/3D arrays, and (q, p) NamedTuples.
• SymmetricMatrix parameter gradient structure: Binding A = ps.A once in each :symmetric method (src/layers/symplectic_attention.jl:141, 150 and src/layers/inverse_layers.jl:131) preserves g.L1.A isa SymmetricMatrix under Zygote.gradient for both MatrixSoftmax() and VectorSoftmax(), Q and P, forward and inverse layers, on 2D and 3D inputs.
• Type stability and allocations: @inferred passes on 1D, 2D, 3D arrays and NamedTuple inputs for forward and inverse layers. In a cold process on (M, N) = (6, 4) (Float64), @allocated is 2704 bytes (2 * σXAX) vs 2720 bytes (σXAX + _custom_transpose(σXAX)); on (M, N, B) = (6, 4, 2), @allocated is 5072 bytes vs 5856 bytes (the extra permutedims(σXAX, (2, 1, 3)) allocation of a 4×4×2 array).
• Test suite, formatting, linting, and layout: Ran layers/symplectic_attention.jl, layers/inverse_layers.jl, layers/symplectic_attention_network_parameters_gradient.jl, and quality/aqua.jl (852/852 passed); JuliaFormatter.format(..., overwrite = false) on all 5 changed .jl files (formatted); fatou lint --force-exclude (0 findings); test-layout.jl --check (0 violations); closingIssuesReferences is empty.

Not checked

• Full Pkg.test() locally: CI ran the full test suite across Julia 1 and Julia min on Ubuntu, macOS, and Windows (plus Metal on macOS-15 and Doctests/Documentation), all green.
• GPU/JLArray parameter gradient pullback for SymplecticAttention: verified in a cold process with JLArrays and allowscalar(false) that the forward pass l(z_jl, ps_jl) succeeds without scalar indexing on 2D and 3D JLArray{Float32} inputs, whereas parameter gradient accumulation Zygote.accum(::SymmetricMatrix{Float32, <:JLArray}, ::SymmetricMatrix{Float32, <:JLArray}) hits pre-existing Base.Broadcast scalar indexing on SymmetricMatrix in GeometricOptimizers.

[michakraus]

Review follow-up (23a0228f)

#	severity	file:line	action	resolution
1	nit	docs/src/layers/symplectic_attention.md:24,26,52,54,60,78	not fixed	Lines 24, 26, 52, 54, 60, 66, 72, and 83–86 are pre-existing text outside the changed hunk (docs/src/layers/symplectic_attention.md:75-82). Inside the changed hunk, \nabla_Z\Sigma(Z) = AZ(\mathrm{softmax}_1(C) + \mathrm{softmax}_1(C)^T) on line 78 is the exact sum of TermI (AZ(\mathrm{softmax}_1(C))^T, line 72) and TermII (A^TZ\mathrm{softmax}_1(C), line 66) for symmetric A, and both standard column-wise \mathrm{softmax} and \mathrm{softmax}_1 share the gradient form A Z (\sigma(C) + \sigma(C)^T). This also addresses the coderabbitai[bot] comment (#discussion_r4454495583) and walkthrough summary (#issuecomment-6603405300).
2	nit	test/layers/inverse_layers.jl:87-94	fixed (23a0228f)	Added SymplecticAttentionQ(6; symmetric = false, activation = VectorSoftmax()) and SymplecticAttentionP(6; symmetric = false, activation = VectorSoftmax()) to @testset "exact layerwise inversion ($T)", completing the (symmetric, activation) grid (test/layers/inverse_layers.jl: 931/931 passed).

[michakraus]

Review follow-up (e0f73284)

#	severity	file:line	action	resolution
1	nit	docs/src/layers/symplectic_attention.md:23-86	fixed (e0f73284)	Addressed coderabbitai[bot] (#discussion_r4454495583 and #issuecomment-6603405300) and the review finding: rewrote the Matrix Softmax and Vector Softmax potentials and gradients to use the standard softmax (\sum \exp(C) denominator and \mathrm{softmax}(C)) matching MatrixSoftmax and VectorSoftmax, kept the miller2023attention one softmax note as an !!! info remark, and fixed the typos "staring point", \sum{m, n}, \exp(C), and "than".
2	nit	test/layers/inverse_layers.jl:87-94	fixed (23a0228f)	Added SymplecticAttentionQ(6; symmetric = false, activation = VectorSoftmax()) and SymplecticAttentionP(6; symmetric = false, activation = VectorSoftmax()) to @testset "exact layerwise inversion ($T)", completing the (symmetric, activation) grid (test/layers/inverse_layers.jl: 931/931 passed).