Make symmetric SymplecticAttention symplectic with VectorSoftmax

CHANGELOG.md

@@ -15,6 +15,10 @@ breaking release).
 
 ## [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 (`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.
+
 ### Changed
 
 - **The `[compat]` lower bounds of `Project.toml` and `test/Project.toml` are the versions the CI Downgrade job resolves when it resolves both projects together at Julia 1.12.** Each old bound named a version that no resolution at the floors reaches:

docs/src/layers/symplectic_attention.md

@@ -20,66 +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\mathrm{softmax}_1(C).
+\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).
 

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})(

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.

test/layers/inverse_layers.jl

@@ -87,7 +87,11 @@ 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()),
+        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)

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

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")