From b869c8e43b9f12c84712471c90a57a4e6a02a12e Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Tue, 3 Nov 2026 12:54:59 +0100 Subject: [PATCH 1/5] Make symmetric SymplecticAttention a gradient shear for VectorSoftmax MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The symmetric method computed the increment 2 A X σ(XᵀAX) for either activation. That is the gradient of the attention potential only when σ is symmetric, which holds for MatrixSoftmax but not for the column-wise VectorSoftmax; there the layer was not symplectic. Compute A X (σ + σᵀ), the gradient for both activations, in the layer and in its inverse, and correct the symmetric formula in the docs page. Co-Authored-By: Claude Opus 5.5 --- docs/src/layers/symplectic_attention.md | 4 +++- src/layers/inverse_layers.jl | 3 ++- src/layers/symplectic_attention.jl | 8 ++++++- test/layers/inverse_layers.jl | 4 +++- test/layers/symplectic_attention.jl | 28 +++++++++++++++++++++++++ 5 files changed, 43 insertions(+), 4 deletions(-) diff --git a/docs/src/layers/symplectic_attention.md b/docs/src/layers/symplectic_attention.md index 9cd96e145..8bd246744 100644 --- a/docs/src/layers/symplectic_attention.md +++ b/docs/src/layers/symplectic_attention.md @@ -75,9 +75,11 @@ The first term is equivalent to: If we again assume that the matrix `A` is a [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix) then the expression simplifies to: ```math -\nabla_Z\Sigma(Z) = AZ\mathrm{softmax}_1(C). +\nabla_Z\Sigma(Z) = AZ\left(\mathrm{softmax}_1(C) + \mathrm{softmax}_1(C)^T\right). ``` +Unlike in the matrix case, the two terms do not merge into one: ``C`` is symmetric, but its column-wise softmax is not. + !!! info Note that we used the "one softmax" here instead of the standard softmax. The one softmax has been shown to be favorable to the standard softmax in some cases. diff --git a/src/layers/inverse_layers.jl b/src/layers/inverse_layers.jl index 1048edba8..f5afb53e1 100644 --- a/src/layers/inverse_layers.jl +++ b/src/layers/inverse_layers.jl @@ -129,7 +129,8 @@ function _attention_displacement( d::InverseSymplecticAttention{M, M, LayerType, :symmetric}, x::AbstractArray, ps::NamedTuple) where {M, LayerType} A = ps.A - _mul(_mul(A, x), 2 * d.activation(_mul(_mul(_custom_transpose(x), A), x))) + σXAX = d.activation(_mul(_mul(_custom_transpose(x), A), x)) + _mul(_mul(A, x), σXAX + _custom_transpose(σXAX)) end function (d::InverseSymplecticAttention{M, M, LayerType})( diff --git a/src/layers/symplectic_attention.jl b/src/layers/symplectic_attention.jl index b184004b6..752f99728 100644 --- a/src/layers/symplectic_attention.jl +++ b/src/layers/symplectic_attention.jl @@ -131,11 +131,17 @@ end # to a plain `Matrix` -- `test/layers/symplectic_attention_network_parameters_gradient.jl` # asserts both halves of that, and the last testset there covers this method. The binding costs no # allocation; it is the number of accesses that decides the structure. +# +# The increment is the gradient `A X (σ + σᵀ)` of the attention potential, and not `2 A X σ`: the +# column-wise `VectorSoftmax` of the symmetric `XᵀAX` is not symmetric. function (d::SymplecticAttention{M, M, LayerType, :symmetric})( z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {M, LayerType, AT} A = ps.A _shift(Val(LayerType), z, - x -> _mul(_mul(A, x), 2 * d.activation(_mul(_mul(_custom_transpose(x), A), x)))) + function (x) + σXAX = d.activation(_mul(_mul(_custom_transpose(x), A), x)) + _mul(_mul(A, x), σXAX + _custom_transpose(σXAX)) + end) end # A single sample, a `(q, p)` pair of vectors, goes through the matrix method as one column. diff --git a/test/layers/inverse_layers.jl b/test/layers/inverse_layers.jl index ad87e0938..5991756f3 100644 --- a/test/layers/inverse_layers.jl +++ b/test/layers/inverse_layers.jl @@ -87,7 +87,9 @@ end SymplecticAttentionQ(6), SymplecticAttentionP(6), SymplecticAttentionQ(6; symmetric = false), - SymplecticAttentionP(6; symmetric = false)) + SymplecticAttentionP(6; symmetric = false), + SymplecticAttentionQ(6; activation = VectorSoftmax()), + SymplecticAttentionP(6; activation = VectorSoftmax())) zv = randn(T, 6) zm = randn(T, 6, 3) zt = randn(T, 6, 3, 2) diff --git a/test/layers/symplectic_attention.jl b/test/layers/symplectic_attention.jl index 13aeb0b0c..467513be3 100644 --- a/test/layers/symplectic_attention.jl +++ b/test/layers/symplectic_attention.jl @@ -1,7 +1,9 @@ using GeometricMachineLearning using GeometricMachineLearning: params using KernelAbstractions +using LinearAlgebra: I using Test +using Zygote: jacobian import Random include("../helpers/eltypes.jl") @@ -29,3 +31,29 @@ Random.seed!(1234) @test !(y ≈ z) end end + +# The Jacobian of the layer on a sequence `(Q, P)` of `L` vectors preserves the canonical form +# that pairs `vec(Q)` with `vec(P)`, for either activation. +function sequence_jacobian(l, ps, Q::AbstractMatrix, P::AbstractMatrix) + n = length(Q) + function f(w) + out = l((q = reshape(w[1:n], size(Q)), p = reshape(w[(n + 1):end], size(P))), ps) + vcat(vec(out.q), vec(out.p)) + end + jacobian(f, vcat(vec(Q), vec(P)))[1] +end + +@testset "symplecticity ($(nameof(typeof(activation))), symmetric = $symmetric)" for activation in ( + MatrixSoftmax(), VectorSoftmax()), + symmetric in (true, false) + + d, L = 2, 4 + Ω = [zeros(d * L, d * L) I; -I zeros(d * L, d * L)] + for l in (SymplecticAttentionQ(2d; symmetric = symmetric, activation = activation), + SymplecticAttentionP(2d; symmetric = symmetric, activation = activation)) + ps = params(NeuralNetwork(Chain(l), CPU(), Float64)).L1 + J = sequence_jacobian(l, ps, randn(d, L), randn(d, L)) + @test J' * Ω * J ≈ Ω atol = 1e-12 + @test !(J ≈ I) + end +end From 525ab3907c732146bd8488d177693b9292f11d25 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Tue, 3 Nov 2026 12:56:20 +0100 Subject: [PATCH 2/5] Add the CHANGELOG entry for the symmetric VectorSoftmax attention fix Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 26980acb6..9b7cf8dc0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,12 @@ breaking release). > alongside the work. Where a release removed exported names the list is given; where it is a > reconstruction of intent, it says so. +## [Unreleased] + +### Fixed + +- **`SymplecticAttentionQ` and `SymplecticAttentionP` with `symmetric = true` are symplectic with `VectorSoftmax`.** The shear increment was `2 A X σ(XᵀAX)`, which is the gradient of the attention potential only when `σ(XᵀAX)` is symmetric. `MatrixSoftmax`, the default, gives a symmetric matrix, so it was correct; the column-wise `VectorSoftmax` does not, and with it the Jacobian of the layer was not symplectic (entries of `JᵀΩJ − Ω` up to about 2e-2 at the default initialisation). The increment is now `A X (σ + σᵀ)`, the gradient for both activations. With `MatrixSoftmax` the output is unchanged up to round-off. With `VectorSoftmax` and `symmetric = true` a trained model keeps its parameters but gives different outputs. `InverseSymplecticAttentionQ` and `InverseSymplecticAttentionP` use the same formula, so they still invert the layer. The `symmetric = false` layers were already correct and are unchanged. The symplectic attention documentation page carries the corrected formula. + ## [0.9.0] — 2026-10-31 ### Removed (breaking) From 8dbb8b0a5ddf10657553c8f417e0d06c304b78ed Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Tue, 3 Nov 2026 13:05:08 +0100 Subject: [PATCH 3/5] Correct the defect size in the CHANGELOG entry and relabel the attention test file Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- test/runtests.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9b7cf8dc0..44f81bbd5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,7 +17,7 @@ breaking release). ### Fixed -- **`SymplecticAttentionQ` and `SymplecticAttentionP` with `symmetric = true` are symplectic with `VectorSoftmax`.** The shear increment was `2 A X σ(XᵀAX)`, which is the gradient of the attention potential only when `σ(XᵀAX)` is symmetric. `MatrixSoftmax`, the default, gives a symmetric matrix, so it was correct; the column-wise `VectorSoftmax` does not, and with it the Jacobian of the layer was not symplectic (entries of `JᵀΩJ − Ω` up to about 2e-2 at the default initialisation). The increment is now `A X (σ + σᵀ)`, the gradient for both activations. With `MatrixSoftmax` the output is unchanged up to round-off. With `VectorSoftmax` and `symmetric = true` a trained model keeps its parameters but gives different outputs. `InverseSymplecticAttentionQ` and `InverseSymplecticAttentionP` use the same formula, so they still invert the layer. The `symmetric = false` layers were already correct and are unchanged. The symplectic attention documentation page carries the corrected formula. +- **`SymplecticAttentionQ` and `SymplecticAttentionP` with `symmetric = true` are symplectic with `VectorSoftmax`.** The shear increment was `2 A X σ(XᵀAX)`, which is the gradient of the attention potential only when `σ(XᵀAX)` is symmetric. `MatrixSoftmax`, the default, gives a symmetric matrix, so it was correct; the column-wise `VectorSoftmax` does not, and with it the Jacobian of the layer was not symplectic (`JᵀΩJ − Ω` is of order one at the default initialisation). The increment is now `A X (σ + σᵀ)`, the gradient for both activations. With `MatrixSoftmax` the output is unchanged up to round-off. With `VectorSoftmax` and `symmetric = true` a trained model keeps its parameters but gives different outputs. `InverseSymplecticAttentionQ` and `InverseSymplecticAttentionP` use the same formula, so they still invert the layer. The `symmetric = false` layers were already correct and are unchanged. The symplectic attention documentation page carries the corrected formula. ## [0.9.0] — 2026-10-31 diff --git a/test/runtests.jl b/test/runtests.jl index 388ea824d..874f56716 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -84,7 +84,7 @@ if "core" in GROUPS @safetestset "Optimizers with structured (non-manifold) weights" include("optimizers/structured_array_parameters.jl") @safetestset "_mul: NetworkParameters gradient structure" include("layers/double_multiplication_network_parameters_gradient.jl") @safetestset "Symplectic attention: NetworkParameters gradient structure" include("layers/symplectic_attention_network_parameters_gradient.jl") - @safetestset "Symplectic attention on a single sample" include("layers/symplectic_attention.jl") + @safetestset "Symplectic attention: single sample and symplecticity" include("layers/symplectic_attention.jl") @safetestset "ZygotePullback returns a NetworkParameters" include("pullbacks/zygote_pullback.jl") @safetestset "AutoEncoderLoss and a loss's NetworkParameters" include("loss/losses.jl") @safetestset "apply_parametric and the concatenation of the system parameters" include("parametric.jl") From 23a0228f6d12e080d6e25c4c296f53464129b807 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Tue, 3 Nov 2026 16:08:51 +0100 Subject: [PATCH 4/5] Test nonsymmetric VectorSoftmax layer inversion --- test/layers/inverse_layers.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/layers/inverse_layers.jl b/test/layers/inverse_layers.jl index 5991756f3..e9eea9d4e 100644 --- a/test/layers/inverse_layers.jl +++ b/test/layers/inverse_layers.jl @@ -89,7 +89,9 @@ end SymplecticAttentionQ(6; symmetric = false), SymplecticAttentionP(6; symmetric = false), SymplecticAttentionQ(6; activation = VectorSoftmax()), - SymplecticAttentionP(6; activation = VectorSoftmax())) + SymplecticAttentionP(6; activation = VectorSoftmax()), + SymplecticAttentionQ(6; symmetric = false, activation = VectorSoftmax()), + SymplecticAttentionP(6; symmetric = false, activation = VectorSoftmax())) zv = randn(T, 6) zm = randn(T, 6, 3) zt = randn(T, 6, 3, 2) From e0f732840742d6f2ebd083eebdbb57732df390a9 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Tue, 3 Nov 2026 17:03:14 +0100 Subject: [PATCH 5/5] Document standard softmax in symplectic attention --- docs/src/layers/symplectic_attention.md | 26 ++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/docs/src/layers/symplectic_attention.md b/docs/src/layers/symplectic_attention.md index 8bd246744..dbe1a4157 100644 --- a/docs/src/layers/symplectic_attention.md +++ b/docs/src/layers/symplectic_attention.md @@ -20,68 +20,68 @@ Its element-wise derivative is: ## Matrix Softmax -Here we take as a staring point the expression: +Here we take as a starting point the expression: ```math -\Sigma(Z) = \mathrm{log}(1 + \exp(\sum_{m,n}C_{mn})). +\Sigma(Z) = \mathrm{log}\left(\sum_{m,n}\exp(C_{mn})\right). ``` Its gradient (with respect to ``Z``) is: ```math -\frac{\partial\Sigma(Z)}{\partial{}Z_{ij}} & = \frac{1}{1 + \sum{m, n}\exp(C_{mn})}\sum_{m'n'}\exp(C_{m'n'})\sum_{\ell}(\delta_{jm'}A_{i\ell}Z_{\ell{}n'} + \delta_{jn'}Z_{\ell{}m'}A_{\ell{}i}) \\ & = \frac{1}{1 + \sum_{m,n}\exp(C_{mn})}\{[AZ\exp.(C)^T]_{ij} + [A^TZ\exp.(C)]_{ij}\}. +\frac{\partial\Sigma(Z)}{\partial{}Z_{ij}} & = \frac{1}{\sum_{m,n}\exp(C_{mn})}\sum_{m'n'}\exp(C_{m'n'})\sum_{\ell}(\delta_{jm'}A_{i\ell}Z_{\ell{}n'} + \delta_{jn'}Z_{\ell{}m'}A_{\ell{}i}) \\ & = \frac{1}{\sum_{m,n}\exp(C_{mn})}\{[AZ\exp.(C)^T]_{ij} + [A^TZ\exp.(C)]_{ij}\}. ``` -Note that if `A` is a [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix) the expression than simplifies to: +Note that if `A` is a [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix) the expression then simplifies to: ```math -\frac{\partial\Sigma(Z)}{\partial{}Z_{ij}} = 2\frac{1}{1 + \sum_{m,n}\exp(C_{mn})}[AZ\exp.(C)^T]_{ij}, +\frac{\partial\Sigma(Z)}{\partial{}Z_{ij}} = 2\frac{1}{\sum_{m,n}\exp(C_{mn})}[AZ\exp.(C)^T]_{ij}, ``` or written in matrix notation: ```math -\nabla_Z\Sigma(Z) = 2\frac{1}{1 + \sum_{m,n}\exp(C_{mn})}AZ\exp.(C). +\nabla_Z\Sigma(Z) = 2\frac{1}{\sum_{m,n}\exp(C_{mn})}AZ\exp.(C). ``` Whether we use a [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix) for ``A`` or not can be set with the keyword `symmetric` in [`SymplecticAttention`](@ref). ## Vector Softmax -Here we take as a staring point the expression: +Here we take as a starting point the expression: ```math -\Sigma(Z) = \sum_{n}\mathrm{log}(1 + \exp(\sum_{m}C_{mn})). +\Sigma(Z) = \sum_{n}\mathrm{log}\left(\sum_{m}\exp(C_{mn})\right). ``` We then get: ```math -\frac{\partial\Sigma(Z)}{\partial{}Z_{ij}} = \sum_n\frac{\exp(C_{jn})}{1 + \sum_m\exp(C_{mn})}[AZ]_{in} + \frac{1}{1 + \sum_m\exp(C_{mj})}[A^TZ\exp(C)]_{ij}. +\frac{\partial\Sigma(Z)}{\partial{}Z_{ij}} = \sum_n\frac{\exp(C_{jn})}{\sum_m\exp(C_{mn})}[AZ]_{in} + \frac{1}{\sum_m\exp(C_{mj})}[A^TZ\exp.(C)]_{ij}. ``` The second term in this expression is equivalent to a *standard attention step*: ```math -\mathrm{TermII:}\qquad A^TZ\mathrm{softmax}_1(C). +\mathrm{TermII:}\qquad A^TZ\mathrm{softmax}(C). ``` The first term is equivalent to: ```math -\mathrm{TermI:}\qquad \sum_n [AZ]_{in}[\mathrm{softmax}_1(C)^T]_{nj} \equiv AZ(\mathrm{softmax}_1(C))^T. +\mathrm{TermI:}\qquad \sum_n [AZ]_{in}[\mathrm{softmax}(C)^T]_{nj} \equiv AZ(\mathrm{softmax}(C))^T. ``` If we again assume that the matrix `A` is a [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix) then the expression simplifies to: ```math -\nabla_Z\Sigma(Z) = AZ\left(\mathrm{softmax}_1(C) + \mathrm{softmax}_1(C)^T\right). +\nabla_Z\Sigma(Z) = AZ\left(\mathrm{softmax}(C) + \mathrm{softmax}(C)^T\right). ``` Unlike in the matrix case, the two terms do not merge into one: ``C`` is symmetric, but its column-wise softmax is not. !!! info - Note that we used the "one softmax" here instead of the standard softmax. The one softmax has been shown to be favorable to the standard softmax in some cases. + Replacing the denominator ``\sum_m\exp(C_{mn})`` with ``1 + \sum_m\exp(C_{mn})`` gives the "one softmax" ``\mathrm{softmax}_1``, which has the same gradient form and is favorable to the standard softmax in some cases. A discussion of the one softmax can be found in [miller2023attention](@cite).