diff --git a/CHANGELOG.md b/CHANGELOG.md index d43011ad3..724347bfa 100644 --- a/CHANGELOG.md +++ b/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: diff --git a/docs/src/layers/symplectic_attention.md b/docs/src/layers/symplectic_attention.md index 9cd96e145..dbe1a4157 100644 --- a/docs/src/layers/symplectic_attention.md +++ b/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). 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..e9eea9d4e 100644 --- a/test/layers/inverse_layers.jl +++ b/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) 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 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")