From 3cac7dafe05e8d80d63bfae4bf91782e1868dd7a Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 30 Oct 2026 23:36:45 +0100 Subject: [PATCH 1/6] Correct the multi-modal losses and the claims of the particle and parity frontends CrossModalCycleLoss replaces its cycle term, which is zero by construction for a weight-tied core with a right-invertible frontend, by the cross-reconstruction error decode(b, z_a) - x_b (beta_cycle becomes beta_cross = 1.0), recombines a pre-split parity pair before the comparison, and raises a DimensionMismatch for x_a and x_b with different numbers of samples. The self-reconstruction forms of CharbonnierRelativeLoss skip the latent term. IsometricRolloutLoss has no default latent Hamiltonian: alpha_H > 0 requires H. The docstrings and two docs pages state the limits of the sorting, slot and parity frontends, and a new testset checks HRedSys with a reference that is not an equilibrium. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 4 +- KNOWN_ISSUES.md | 6 +- docs/src/layers/stiefel_layer.md | 11 +- docs/src/layers/symmetric_particle_layers.md | 4 +- .../multimodal_symplectic_autoencoder.jl | 3 +- .../symmetric_symplectic_autoencoder.jl | 1 + src/layers/bilinear_stiefel_layer.jl | 5 + src/layers/canonical_slot_layer.jl | 11 +- src/layers/canonical_sort_layer.jl | 3 + src/layers/low_rank_stiefel_layer.jl | 5 + src/loss/multimodal_losses.jl | 120 ++++++-- .../multimodal_symplectic_autoencoder.jl | 263 ++++++++++++++++-- .../reduced_system/reference_shift_and_vjp.jl | 67 +++++ 13 files changed, 442 insertions(+), 61 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 583baa706..1d59f74f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -150,10 +150,12 @@ breaking release). - **`BilinearStiefelLayer`, `TransposeBilinearStiefelLayer`, `LowRankStiefelLayer`, `DecodeLowRankStiefelLayer`, `FastDLRATraceLoss`, and `solve!` on `BilinearStiefelLayer` add two-sided velocity-parity Stiefel frontends, a gauge-invariant dynamical low-rank Stiefel frontend and decoder, a fast low-rank trace loss, and Tucker-2 HOSVD initialization.** The five new types are exported. `BilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)` (`Nx >= 1`, even `Nv >= 2`, `1 <= mx <= Nx`, `1 <= mv <= Nv ÷ 2`) allocates `(Φx = StiefelManifold(...), Φv = StiefelManifold(...))` with ``\Phi_x \in \mathrm{St}(m_x, N_x)`` and ``\Phi_v \in \mathrm{St}(m_v, N_v / 2)``, splits an input array ``F`` of shape `(Nx, Nv)` or `(Nx, Nv, B)` into velocity-even ``F^+ = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} + F_{:, N_v:-1:N_v/2+1})`` and velocity-odd ``F^- = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} - F_{:, N_v:-1:N_v/2+1})`` components (or accepts a pre-split `(q, p)` `NamedTuple` of shape `(Nx, Nv ÷ 2)` or `(Nx, Nv ÷ 2, B)`), applies optional Cholesky mass factors `(Rx, Rv)` when `metric_factors` is supplied, and projects ``Q_{\mathrm{red}} = \Phi_x^T F^+ \Phi_v`` and ``P_{\mathrm{red}} = \Phi_x^T F^- \Phi_v`` to return vectorized `(q = vec(Q_red), p = vec(P_red))` of half-dimension ``n = m_x m_v``; `TransposeBilinearStiefelLayer` (`inverse_layer(l)`) is its exact right-inverse. `LowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)` takes a DLRA `NamedTuple` `(X = X, S = S, V = V)` of arbitrary rank ``r_{\mathrm{in}} \ge 1`` and evaluates ``Q_{\mathrm{red}} = (\Phi_x^T X) S ((V^+)^T \Phi_v)`` and ``P_{\mathrm{red}} = (\Phi_x^T X) S ((V^-)^T \Phi_v)`` in ``O(r_{\mathrm{in}}(m_x N_x + m_v N_v))`` operations without forming ``X S V^T``, matching `BilinearStiefelLayer` on ``F = X S V^T`` identically and remaining invariant under ``O(r_{\mathrm{in}}) \times O(r_{\mathrm{in}})`` gauge transformations, while `DecodeLowRankStiefelLayer` (`inverse_layer(l)`) computes the truncated SVD of ``[Q_{\mathrm{red}}\;\; P_{\mathrm{red}}]`` of rank ``r_{\mathrm{eff}} = \min(r_{\mathrm{out}}, m_x, 2 m_v)`` and QR re-orthonormalizes the lifted factors to return `(X = X_hat, S = S_hat, V = V_hat)` on the Stiefel manifold. `FastDLRATraceLoss()` evaluates ``\|X S V^T - \hat{X} \hat{S} \hat{V}^T\|_F^2`` in ``O(r^2(N_x + N_v))`` operations via Gram-matrix traces clamped at `zero(T)`. `solve!(nn, F_snapshots)` and `solve!(l, ps, F_snapshots)` initialize `Φx` and `Φv` on a `BilinearStiefelLayer` from the top-``m_x`` spatial and top-``m_v`` parity-velocity left singular vectors of `F_snapshots`. The tests are `test/layers/bilinear_and_low_rank_stiefel.jl`. -- **`MultiModalSymplecticAutoencoder`, `encode`, `decode`, `translate`, `CharbonnierRelativeLoss`, `IsometricRolloutLoss`, and `CrossModalCycleLoss` add multi-modal symplectic compression and multi-modal training losses, and `HRedSys` gains reference-state shifting, reverse-mode VJP vector fields, and preallocated work buffers.** All seven new names are exported. `MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression` combines a non-empty `NamedTuple` of modality frontends (individual layers or `Chain`s of symplectic particle layers) whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` (with `modality` a `Symbol` or `Val(modality)`) applies `frontends[modality]` followed by `encoder(core)`, `decode(nn, modality, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])`, and `translate(nn, from, to, x)` computes `decode(nn, to, encode(nn, from, x))`, raising an `ArgumentError` on an empty or dimension-mismatched `frontends` or an unknown modality symbol. `CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)` evaluates the exact-zero Charbonnier relative reconstruction loss ``(\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon) / \sqrt{\|x\|_F^2 + \varepsilon^2} + \lambda_{\mathrm{lat}} \|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2 / (\|x\|_F^2 + \varepsilon^2)`` (supporting low-rank `(X, S, V)` inputs via `FastDLRATraceLoss`), `IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.1, H = nothing)` evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty, and `CrossModalCycleLoss(mod_a, mod_b; β_align = 1.0, β_cycle = 0.1)` evaluates paired latent alignment ``\|z_a - z_b\|_F^2`` plus cycle consistency ``\|\mathrm{encode}(nn, b, \mathrm{decode}(nn, b, z_a)) - z_a\|_F^2``. `IsometricRolloutLoss` raises an `ArgumentError` naming both element types when the data, the output of `integrator` or `H` has another element type than the network's parameters, and `CrossModalCycleLoss` does so for `x_a` or `x_b`. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (converted to the state element type, shifting the decoder lift by ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` and subtracting ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` before `rs.encoder` in `integrate_reduced_system`) and `pullback_mode::Symbol = :jacobian | :vjp`, where `:vjp` evaluates ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` in one reverse-mode pullback `Zygote.pullback(decoder, ξ)` and both modes use the preallocated buffers of `cache` (`ξ_buf`, `v_intermediate`, `f_intermediate`, and `jac_result` for `:jacobian`, with a per-element-type `IdDict{DataType, Any}` of caches for `ForwardDiff.Dual` Newton states). The default `cache` of `build_v_reduced` and `build_f_reduced` has the element type of the decoder's parameters, and the default `cache` of `HRedSys` has the element type of `ics.q`; a state of another floating-point element type than the cache's raises an `ArgumentError`. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. +- **`MultiModalSymplecticAutoencoder`, `encode`, `decode`, `translate`, `CharbonnierRelativeLoss`, `IsometricRolloutLoss`, and `CrossModalCycleLoss` add multi-modal symplectic compression and multi-modal training losses, and `HRedSys` gains reference-state shifting, reverse-mode VJP vector fields, and preallocated work buffers.** All seven new names are exported. `MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression` combines a non-empty `NamedTuple` of modality frontends (individual layers or `Chain`s of symplectic particle layers) whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` (with `modality` a `Symbol` or `Val(modality)`) applies `frontends[modality]` followed by `encoder(core)`, `decode(nn, modality, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])`, and `translate(nn, from, to, x)` computes `decode(nn, to, encode(nn, from, x))`, raising an `ArgumentError` on an empty or dimension-mismatched `frontends` or an unknown modality symbol. `CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)` evaluates, for an external reconstruction `x̂`, the exact-zero Charbonnier relative reconstruction loss ``(\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon) / \sqrt{\|x\|_F^2 + \varepsilon^2} + \lambda_{\mathrm{lat}} \|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2 / (\|x\|_F^2 + \varepsilon^2)``, and without `x̂` the reconstruction term alone for ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` (supporting low-rank `(X, S, V)` inputs via `FastDLRATraceLoss`), `IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.1, H = nothing)` evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty, with an `ArgumentError` when `α_H > 0` and `H = nothing`, and `CrossModalCycleLoss(mod_a, mod_b; β_align = 1.0, β_cross = 1.0)` evaluates paired latent alignment ``\|z_a - z_b\|_F^2`` plus the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``. `IsometricRolloutLoss` raises an `ArgumentError` naming both element types when the data, the output of `integrator` or `H` has another element type than the network's parameters, and `CrossModalCycleLoss` does so for `x_a` or `x_b`. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (converted to the state element type, shifting the decoder lift by ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` and subtracting ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` before `rs.encoder` in `integrate_reduced_system`) and `pullback_mode::Symbol = :jacobian | :vjp`, where `:vjp` evaluates ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` in one reverse-mode pullback `Zygote.pullback(decoder, ξ)` and both modes use the preallocated buffers of `cache` (`ξ_buf`, `v_intermediate`, `f_intermediate`, and `jac_result` for `:jacobian`, with a per-element-type `IdDict{DataType, Any}` of caches for `ForwardDiff.Dual` Newton states). The default `cache` of `build_v_reduced` and `build_f_reduced` has the element type of the decoder's parameters, and the default `cache` of `HRedSys` has the element type of `ics.q`; a state of another floating-point element type than the cache's raises an `ArgumentError`. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. ### Changed +- **The multi-modal losses of #370 lose the terms that are zero by construction, `IsometricRolloutLoss` has no default latent Hamiltonian, and the docstrings of the particle and parity frontends state their limits.** No release has these names, so nothing is deprecated. `CrossModalCycleLoss` replaces its cycle term ``\|\mathrm{encode}(nn, b, \mathrm{decode}(nn, b, z_a)) - z_a\|_F^2``, which is zero for a weight-tied core with a right-invertible frontend, with the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``, which depends on the data `x_b`; the keyword `β_cycle = 0.1` becomes `β_cross = 1.0`, a passed `β_cycle` raises Julia's `MethodError`, and a negative or `NaN` `β_cross` raises an `ArgumentError`. `x_a` and `x_b` whose latent states have different sizes, such as one sample against a batch, raise a `DimensionMismatch`. The self-reconstruction forms of `CharbonnierRelativeLoss`, those without an external `x̂`, skip the latent term, which is zero there for a weight-tied core with a right-invertible frontend, so `λ_lat` acts only in the form with `x̂`. `IsometricRolloutLoss(integrator)` raises an `ArgumentError`, because `α_H > 0` requires a latent Hamiltonian `H` and the fixed ``\frac{1}{2}\|\hat{z}\|_2^2`` of #370 is not the system's energy; with `α_H = 0` it needs none and calls none. The cross-reconstruction term, the self-reconstruction form of `CharbonnierRelativeLoss` and the `α_rec` term of `IsometricRolloutLoss` compare a decoded `BilinearStiefelLayer` field with a pre-split pair `(q = F⁺, p = F⁻)` after recombining the pair into the full field. The docstrings of `CanonicalSortLayer`, `CanonicalSlotLayer`, `InverseCanonicalSlotLayer`, `SymmetricSymplecticAutoencoder`, `BilinearStiefelLayer`, `LowRankStiefelLayer`, `MultiModalSymplecticAutoencoder` and the three losses, and the pages `docs/src/layers/symmetric_particle_layers.md` and `docs/src/layers/stiefel_layer.md`, state that sorting is discontinuous where two particles cross and suits systems that keep their particle order, that the slot layer is a linear projection of sorted states for a fixed particle count, and that the `(q, p)` of the parity frontends is canonical only for the structure that they declare on the parity coordinates, not for the Vlasov bracket. A new testset of `test/reduced_system/reference_shift_and_vjp.jl` checks `HRedSys` with a reference that is not an equilibrium against the unshifted system in shifted coordinates. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. + - **`SymplecticAttentionQ` and `SymplecticAttentionP` accept a single sample.** A `(q, p)` `NamedTuple` of vectors, and a plain vector, are reshaped to one column, passed through the matrix method and returned as vectors, so the forward layer accepts the same inputs as its inverse. On `main` a `NamedTuple` of vectors and a plain vector raised `MethodError: no method matching _custom_transpose(::Vector{Float64})`. The test is `test/layers/symplectic_attention.jl`. - **`SymplecticAutoencoder` builds with one encoder or decoder block, and names the limit when there are too many.** `compute_iterations_for_symplectic_system`, which `SymplecticAutoencoder` and `WeightTiedSymplecticAutoencoder` share, divided by `n_blocks - 1` and raised `DivideError` for `n_blocks = 1`; it now returns `[reduced_dim, full_dim]`, one block. A block count whose spacing `(full_dim - reduced_dim) ÷ 2 ÷ (n_blocks - 1)` rounds to zero raised Base's `ArgumentError: step cannot be zero`, and now raises an `ArgumentError` that names the largest allowed count: `SymplecticAutoencoder(8, 4; n_decoder_blocks = 4)` gives "n_blocks (4) must be at most 3 for full_dim (8) and reduced_dim (4).", and a block count below 1 gives "n_blocks (0) must be at least 1.". As before, the `SymplecticAutoencoder` constructor accepts these block counts, and the error is raised when the network is built. The tests are `test/architectures/symplectic_autoencoder_tests.jl` and `test/architectures/weight_tied_symplectic_autoencoder.jl`. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 975a238e4..c9dfd43bf 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -78,15 +78,15 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. ### K41 · `CrossModalCycleLoss(nn, x_a, x_b)` with a mixed-component `NamedTuple` raises a `MethodError`, not the documented `ArgumentError`. -- **location:** `src/loss/multimodal_losses.jl:323` +- **location:** `src/loss/multimodal_losses.jl:385` - **kind:** defect - **found:** #370 - **evidence:** - The three `(nn, x_a, x_b)` methods (lines 323, 329 and 335) annotate `x_a` and `x_b` with + The three `(nn, x_a, x_b)` methods (lines 385, 391 and 397) annotate `x_a` and `x_b` with `ArrayOrNamedTuple`, and AbstractNeuralNetworks defines that as `Union{AbstractArray{T}, NamedTupleOfArrays{T}}`, with one element type `T` for every component. A `NamedTuple` whose `q` and `p` differ in element type matches no `T`, so dispatch fails before - `_check_eltype` (line 204) runs. In a fresh `julia --startup-file=no --project=.` process: + `_check_eltype` (line 254) runs. In a fresh `julia --startup-file=no --project=.` process: ```julia T = typeof((q = zeros(Float32, 2, 3), p = zeros(Float64, 2, 3))) diff --git a/docs/src/layers/stiefel_layer.md b/docs/src/layers/stiefel_layer.md index 3fa3670f1..8ea9cbf71 100644 --- a/docs/src/layers/stiefel_layer.md +++ b/docs/src/layers/stiefel_layer.md @@ -1,9 +1,14 @@ # Bilinear and Low-Rank Stiefel Layers `BilinearStiefelLayer`, `TransposeBilinearStiefelLayer`, `LowRankStiefelLayer`, -`DecodeLowRankStiefelLayer`, and `FastDLRATraceLoss` provide canonical velocity-parity -Stiefel frontends for two-sided phase-space grid fields ``F \in \mathbb{R}^{N_x \times N_v}`` -and dynamical low-rank approximation (DLRA) states ``(X, S, V)``. +`DecodeLowRankStiefelLayer`, and `FastDLRATraceLoss` provide velocity-parity Stiefel +frontends for two-sided phase-space grid fields ``F \in \mathbb{R}^{N_x \times N_v}`` and +dynamical low-rank approximation (DLRA) states ``(X, S, V)``. + +Their ``(q, p)`` is canonical only for the structure ``\mathbb{J}`` that they declare on the +parity coordinates. A frontend that is linear in ``f`` is not a Poisson map for the Vlasov +bracket: with ``L^2``-orthonormal even and odd bases, the bracket linearised at ``f_0`` gives a +singular ``\{Q, P\}``, not the identity. ## Bilinear Grid and Spline Stiefel Layers diff --git a/docs/src/layers/symmetric_particle_layers.md b/docs/src/layers/symmetric_particle_layers.md index 765da2900..eeb87887f 100644 --- a/docs/src/layers/symmetric_particle_layers.md +++ b/docs/src/layers/symmetric_particle_layers.md @@ -11,7 +11,7 @@ P_\sigma^T \mathbb{J}_{2Md} P_\sigma = \mathbb{J}_{2Md}, \qquad P_\sigma = \begin{bmatrix} \Pi_\sigma \otimes I_d & \mathbb{O} \\ \mathbb{O} & \Pi_\sigma \otimes I_d \end{bmatrix}. ``` -A standard [`PSDLayer`](@ref), [`GradientLayerQ`](@ref), or [`GradientLayerP`](@ref) breaks ``S_M`` symmetry unless the particle state is canonicalized or processed by ``S_M``-equivariant symplectic layers. `GeometricMachineLearning` provides three particle frontends and the composite architecture [`SymmetricSymplecticAutoencoder`](@ref). The sorting and slot frontends are ``S_M``-invariant; the pairwise shears are ``S_M``-equivariant. +A standard [`PSDLayer`](@ref), [`GradientLayerQ`](@ref), or [`GradientLayerP`](@ref) breaks ``S_M`` symmetry unless the particle state is canonicalized or processed by ``S_M``-equivariant symplectic layers. `GeometricMachineLearning` provides three particle frontends and the composite architecture [`SymmetricSymplecticAutoencoder`](@ref). The sorting and slot frontends are invariant under a relabelling of particles with distinct coordinates along the sorting axis; the pairwise shears are ``S_M``-equivariant. Sorting is discontinuous in ``p`` where two particles cross along the sorting axis, so the Jacobian identity of a sorting frontend holds only away from those crossings. The sorting and slot frontends suit systems that keep their particle order, not a collisionless plasma, whose sorted snapshots are not low-rank. ## Canonical Weyl-Chamber Sorting @@ -36,7 +36,7 @@ InversePairwiseSymplecticLayerP ## Canonical Slot Reduction -[`CanonicalSlotLayer`](@ref) sorts the particles along coordinate row 1 of ``q`` with [`CanonicalSortLayer`](@ref), separates the center-of-mass mode ``(Q_{\mathrm{cm}}, P_{\mathrm{cm}})`` from the ``M - 1`` zero-sum relative modes via the orthogonal Helmert / Jacobi transformation, sorts the ``M - 1`` relative modes canonically by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2``, and projects onto ``N_{\mathrm{slots}} - 1`` relative slots via a learned Stiefel matrix ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)``. Its analytic decoder lift is [`InverseCanonicalSlotLayer`](@ref). +[`CanonicalSlotLayer`](@ref) sorts the particles along coordinate row 1 of ``q`` with [`CanonicalSortLayer`](@ref), separates the center-of-mass mode ``(Q_{\mathrm{cm}}, P_{\mathrm{cm}})`` from the ``M - 1`` zero-sum relative modes via the orthogonal Helmert / Jacobi transformation, sorts the ``M - 1`` relative modes by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2``, and projects onto ``N_{\mathrm{slots}} - 1`` relative slots via a learned Stiefel matrix ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)``. Its analytic decoder lift is [`InverseCanonicalSlotLayer`](@ref). The particles are already sorted, so the second sort removes no symmetry, and it makes the output discontinuous where two relative-mode norms cross. The layer is a linear projection of sorted states for a fixed particle count ``M``; it takes no variable particle count and has no Löwdin-orthogonalised affinity slots. ```@docs CanonicalSlotLayer diff --git a/src/architectures/multimodal_symplectic_autoencoder.jl b/src/architectures/multimodal_symplectic_autoencoder.jl index 989a34c5d..62791277f 100644 --- a/src/architectures/multimodal_symplectic_autoencoder.jl +++ b/src/architectures/multimodal_symplectic_autoencoder.jl @@ -29,7 +29,8 @@ ChainRulesCore.@non_differentiable _check_modality(::Any, ::Any) Make an instance of a multi-modal symplectic autoencoder mapping a non-empty `NamedTuple` of modality frontends into a shared canonical latent space ``\xi = (\tilde{q}, \tilde{p}) \in \mathbb{R}^{2n}`` via `core`. -Each entry of `frontends` maps a modality symbol (such as `:particles`, `:spline`, `:grid`, or `:dlra`) to a symplectic or Poisson frontend layer whose canonical output dimension equals `core.full_dim = 2N`: +Each entry of `frontends` maps a modality symbol (such as `:particles`, `:spline`, `:grid`, or `:dlra`) to a frontend layer whose output dimension equals `core.full_dim = 2N`. +The latent state is canonical for the structure that each frontend declares; for the two parity frontends, [`BilinearStiefelLayer`](@ref) and [`LowRankStiefelLayer`](@ref), that structure is the ``\mathbb{J}_{2N}`` that they declare on the parity coordinates, not the Vlasov bracket. The frontends are: - [`CanonicalSortLayer`](@ref), [`CanonicalSlotLayer`](@ref), [`PairwiseSymplecticLayerQ`](@ref), [`PairwiseSymplecticLayerP`](@ref), or a `Chain` of particle layers with output dimension ``2N``; - [`BilinearStiefelLayer`](@ref)`(Nx, Nv, Mx, Mv)` with ``2 M_x M_v = 2N`` (with or without `metric_factors` for spline vs grid modalities); - [`LowRankStiefelLayer`](@ref)`(Nx, Nv, Mx, Mv; r_out = ...)` with ``2 M_x M_v = 2N``. diff --git a/src/architectures/symmetric_symplectic_autoencoder.jl b/src/architectures/symmetric_symplectic_autoencoder.jl index 48e1f459f..ca68b5bbe 100644 --- a/src/architectures/symmetric_symplectic_autoencoder.jl +++ b/src/architectures/symmetric_symplectic_autoencoder.jl @@ -6,6 +6,7 @@ Make an instance of a symplectic autoencoder with a particle frontend for `M` pa `SymmetricSymplecticAutoencoder` combines a particle frontend (`frontend = :sort`, `:pairwise`, or `:slot`) with the weight-tied symplectic core [`WeightTiedSymplecticAutoencoder`](@ref) so that [`encoder`](@ref)`(nn)` and [`decoder`](@ref)`(nn)` share a single parameter tree. With `frontend = :sort` or `frontend = :slot`, the latent state is invariant under a relabelling of particles with distinct coordinates along the sorting axis (`sort_axis` for `:sort`, coordinate 1 for `:slot`). With `frontend = :pairwise`, only the frontend is permutation-equivariant: the core is not, so the latent state is not invariant. +The sort of `:sort` and `:slot` is discontinuous in `p` where two particles cross along the sorting axis, so the Jacobian identity of the frontend holds only away from those crossings. `:sort` and `:slot` suit systems that keep their particle order, not a collisionless plasma, whose sorted snapshots are not low-rank. The latent right-inversion ``\mathrm{encoder}(nn)(\mathrm{decoder}(nn)(\xi)) = \xi`` holds to machine precision for every ``\xi \in \mathbb{R}^{2n}`` with `frontend = :pairwise`. With `frontend = :sort` and `frontend = :slot`, it holds for ``\xi`` whose decoded state lies in the canonical chamber, where the sorting permutations on the decoded state are the identity; the decoder does not map into this chamber in general. diff --git a/src/layers/bilinear_stiefel_layer.jl b/src/layers/bilinear_stiefel_layer.jl index f3db0e6fe..d60bb5dbb 100644 --- a/src/layers/bilinear_stiefel_layer.jl +++ b/src/layers/bilinear_stiefel_layer.jl @@ -30,6 +30,11 @@ When `metric_factors = (Rx, Rv)` is supplied with Cholesky or square-root mass f `Rx` of size `(Nx, Nx)` and `Rv` of size `(Nv ÷ 2, Nv ÷ 2)` or `(Nv, Nv)`, the field is mapped into the Euclidean metric by `(Rx, Rv)` prior to Stiefel projection. +The coordinates ``(q, p)`` are canonical only for the structure ``\mathbb{J}_{2n}`` that the +layer and [`LowRankStiefelLayer`](@ref) declare on the parity coordinates. A frontend that is +linear in ``F`` is not a Poisson map for the Vlasov bracket: with ``L^2``-orthonormal even and +odd bases, the bracket linearised at ``f_0`` gives a singular ``\{Q, P\}``, not the identity. + See also [`TransposeBilinearStiefelLayer`](@ref), [`LowRankStiefelLayer`](@ref), and [`inverse_layer`](@ref). """ diff --git a/src/layers/canonical_slot_layer.jl b/src/layers/canonical_slot_layer.jl index fe32e6f9e..e6261498a 100644 --- a/src/layers/canonical_slot_layer.jl +++ b/src/layers/canonical_slot_layer.jl @@ -41,7 +41,7 @@ end CanonicalSlotLayer{N, NOut, M, NSlots}() CanonicalSlotLayer(M::Integer, d::Integer, NSlots::Integer) -Make an instance of a canonical slot reduction layer mapping `M` particles in ``d`` spatial dimensions to `NSlots` collective canonical slots (`M >= 2`, `1 <= NSlots <= M`). +Make an instance of a canonical slot reduction layer mapping `M` particles in ``d`` spatial dimensions to `NSlots` canonical slots (`M >= 2`, `1 <= NSlots <= M`). The first two type parameters are the phase-space dimensions `N = 2Md` of the particles and `NOut = 2 NSlots d` of the slots, which `input_dimension` and `output_dimension` return. The layer first sorts the particles along coordinate row 1 of `q` with [`CanonicalSortLayer`](@ref), so that its output is invariant under a relabelling of particles with distinct first coordinates. It then applies the orthogonal Helmert / Jacobi transformation separating the center-of-mass mode @@ -49,7 +49,12 @@ The layer first sorts the particles along coordinate row 1 of `q` with [`Canonic ```math (Q_{\mathrm{cm}}, P_{\mathrm{cm}}) = \frac{1}{\sqrt{M}}\sum_{i=1}^M (q_i, p_i) ``` -from the ``M - 1`` zero-sum relative modes ``(q_{\mathrm{rel}}, p_{\mathrm{rel}}) \in \mathbb{R}^{d \times (M - 1)}``, sorts the ``M - 1`` relative modes canonically by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2`` (`sortperm` with `rev = true, alg = MergeSort`, applying the identical permutation ``\sigma_{\mathrm{rel}} \in S_{M - 1}`` to both ``q_{\mathrm{rel}}`` and ``p_{\mathrm{rel}}``), and projects the sorted relative modes onto ``N_{\mathrm{slots}} - 1`` relative slots via a learned [`StiefelManifold`](@extref GeometricOptimizers GeometricOptimizers.StiefelManifold) parameter ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)`` while keeping ``(Q_{\mathrm{cm}}, P_{\mathrm{cm}})`` intact as slot 1. +from the ``M - 1`` zero-sum relative modes ``(q_{\mathrm{rel}}, p_{\mathrm{rel}}) \in \mathbb{R}^{d \times (M - 1)}``, sorts the ``M - 1`` relative modes by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2`` (`sortperm` with `rev = true, alg = MergeSort`, applying the identical permutation ``\sigma_{\mathrm{rel}} \in S_{M - 1}`` to both ``q_{\mathrm{rel}}`` and ``p_{\mathrm{rel}}``), and projects the sorted relative modes onto ``N_{\mathrm{slots}} - 1`` relative slots via a learned [`StiefelManifold`](@extref GeometricOptimizers GeometricOptimizers.StiefelManifold) parameter ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)`` while keeping ``(Q_{\mathrm{cm}}, P_{\mathrm{cm}})`` intact as slot 1. +The particles are already sorted, so this second sort removes no symmetry, and it makes the output discontinuous where two relative-mode norms cross. + +The layer is a linear projection of sorted states for a fixed particle count `M`; it takes no variable particle count and has no Löwdin-orthogonalised affinity slots. +The sort of the particles is discontinuous in `p` where two particles cross along coordinate row 1, so the identities of the layer's Jacobian hold only away from those crossings. +The layer suits systems that keep their particle order, not a collisionless plasma, whose sorted snapshots are not low-rank. """ struct CanonicalSlotLayer{N, NOut, M, NSlots} <: AbstractExplicitLayer{N, NOut} function CanonicalSlotLayer{N, NOut, M, NSlots}() where {N, NOut, M, NSlots} @@ -181,7 +186,7 @@ end InverseCanonicalSlotLayer{N, NOut, M, NSlots}() InverseCanonicalSlotLayer(M::Integer, d::Integer, NSlots::Integer) -Analytic symplectic decoder lift (`inverse_layer(CanonicalSlotLayer(M, d, NSlots))`) mapping `NSlots` collective canonical slots back to `M` particles via ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)`` and the inverse Helmert matrix. +Analytic symplectic decoder lift (`inverse_layer(CanonicalSlotLayer(M, d, NSlots))`) mapping `NSlots` canonical slots back to `M` particles via ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)`` and the inverse Helmert matrix. It takes the type parameters of [`CanonicalSlotLayer`](@ref), so its input dimension is `NOut = 2 NSlots d` and its output dimension is `N = 2Md`. For any ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)``, the decoder lift satisfies ``(\nabla \Psi_{\mathrm{slot}}^{\mathrm{dec}})^T \mathbb{J}_{2Md} \nabla \Psi_{\mathrm{slot}}^{\mathrm{dec}} = \mathbb{J}_{2 N_{\mathrm{slots}} d}`` and preserves ``(Q_{\mathrm{cm}}, P_{\mathrm{cm}})``. diff --git a/src/layers/canonical_sort_layer.jl b/src/layers/canonical_sort_layer.jl index 302ad14db..ba089d52c 100644 --- a/src/layers/canonical_sort_layer.jl +++ b/src/layers/canonical_sort_layer.jl @@ -34,6 +34,9 @@ Make an instance of a cotangent-lifted canonical Weyl-chamber particle sorting l The first type parameter is the phase-space dimension `N = 2Md`, which `input_dimension` and `output_dimension` return. Given a `(q, p)` `NamedTuple` with `q, p` of shape `(d, M)` (single snapshot) or `(d, M, B)` (batch of ``B`` snapshots), `CanonicalSortLayer{N, M, Axis}()` sorts the `M` particles along coordinate row `Axis` of `q` using a stable merge sort (`MergeSort`), applies the **identical** permutation ``\sigma_b \in S_M`` to both `q` and `p` in each batch sample ``b``, and returns `(q = q_sorted, p = p_sorted, perm = σ)`. On a flattened phase-space vector of length ``2Md`` or snapshot matrix of shape ``(2Md, B)``, it reshapes each half to `(d, M)` or `(d, M, B)` and returns the sorted flattened state. + +The sort is discontinuous in `p` where two particles cross along the sorting axis, because the permutation changes there, so the Jacobian identity ``(\nabla\Psi)^T \mathbb{J}_{2Md} \nabla\Psi = \mathbb{J}_{2Md}`` holds only away from those crossings. +Sorting suits systems that keep their particle order, not a collisionless plasma, whose sorted snapshots are not low-rank. """ struct CanonicalSortLayer{N, M, Axis} <: AbstractExplicitLayer{N, N} function CanonicalSortLayer{N, M, Axis}() where {N, M, Axis} diff --git a/src/layers/low_rank_stiefel_layer.jl b/src/layers/low_rank_stiefel_layer.jl index 587918c71..4638252da 100644 --- a/src/layers/low_rank_stiefel_layer.jl +++ b/src/layers/low_rank_stiefel_layer.jl @@ -24,6 +24,11 @@ which produces the exact same canonical coordinates `(q = vec(Q_red), p = vec(P_ orthogonal gauge transformations ``(X, S, V) \mapsto (X O_x, O_x^T S O_v, V O_v)`` for all ``O_x, O_v \in O(r_{\mathrm{in}})``. +The coordinates ``(q, p)`` are canonical only for the structure ``\mathbb{J}_{2n}`` that the layer and +[`BilinearStiefelLayer`](@ref) declare on the parity coordinates. A frontend that is linear in +``f = X S V^T`` is not a Poisson map for the Vlasov bracket: with ``L^2``-orthonormal even and +odd bases, the bracket linearised at ``f_0`` gives a singular ``\{Q, P\}``, not the identity. + See also [`DecodeLowRankStiefelLayer`](@ref), [`BilinearStiefelLayer`](@ref), [`FastDLRATraceLoss`](@ref), and [`inverse_layer`](@ref). """ diff --git a/src/loss/multimodal_losses.jl b/src/loss/multimodal_losses.jl index fac169b3f..c9d8288ce 100644 --- a/src/loss/multimodal_losses.jl +++ b/src/loss/multimodal_losses.jl @@ -51,18 +51,42 @@ function _frob_sq_diff(x̂::NamedTuple{(:X, :S, :V)}, x::NamedTuple{(:X, :S, :V) FastDLRATraceLoss()(x̂, x) end +# The data that a decoded state of a frontend is compared with, in the form that `decode` +# returns for it. `decode` through a `BilinearStiefelLayer` returns the full field, so a pre-split +# pair `(q = F⁺, p = F⁻)` is recombined into that field; every other input is compared as it is. +function _reconstruction_target(::BilinearStiefelLayer, x::NamedTuple{(:q, :p)}) + _parity_recombine(x.q, x.p) +end +_reconstruction_target(::Any, x) = x +function _reconstruction_target(model::MultiModalSymplecticLayer, modality::Symbol, x) + _reconstruction_target(getfield(model.frontends, modality), x) +end + @doc raw""" CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2) -Smooth Charbonnier relative reconstruction and latent consistency loss: +Smooth Charbonnier relative reconstruction loss, with a latent consistency term when the reconstruction ``\hat{x}`` is passed. + +Called as `loss(nn, m, x̂, x)` with an external ``\hat{x}`` of the same form as ``x``, it evaluates ```math \mathcal{L}_{\mathrm{Charb}}(\hat{x}, x) = \frac{\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon}{\sqrt{\|x\|_F^2 + \varepsilon^2}} -+ \lambda_{\mathrm{lat}}\,\frac{\|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2}{\|x\|_F^2 + \varepsilon^2}, ++ \lambda_{\mathrm{lat}}\,\frac{\|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2}{\|x\|_F^2 + \varepsilon^2}. +``` + +Called as `loss(nn, m, x)` (or `loss(nn, x)` for the first modality), it evaluates the reconstruction term alone, + +```math +\mathcal{L}_{\mathrm{Charb}}(x) = +\frac{\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon}{\sqrt{\|x\|_F^2 + \varepsilon^2}}, +\qquad \hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x)), ``` -where ``\varepsilon > 0`` (default `1e-3`) regularizes the Frobenius norm at ``\hat{x} = x`` and ``x = 0``, and ``\lambda_{\mathrm{lat}} \ge 0`` (default `1e-2`) weights the latent re-encoding consistency penalty. +and compares ``\hat{x}`` with ``x`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality is recombined into the full field. +In this form the latent term is zero for a weight-tied core with a right-invertible frontend, because then ``\mathrm{encode}(nn, m, \hat{x}) = \mathrm{encode}(nn, m, x)``, so the method skips it and `λ_lat` has no effect. For any other core, pass ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` to the form with an external ``\hat{x}`` to get the latent term. + +Here ``\varepsilon > 0`` (default `1e-3`) regularizes the Frobenius norm at ``\hat{x} = x`` and ``x = 0``, and ``\lambda_{\mathrm{lat}} \ge 0`` (default `1e-2`) weights the latent re-encoding consistency penalty. When ``x = 0`` and ``\hat{x} = 0``, the loss returns exact `zero(T)` (not `NaN`), and its gradient at ``\hat{x} = x`` is finite and zero (`0 / ε = 0`). """ @@ -92,6 +116,15 @@ function _charbonnier_eval( (term_rec + term_lat)::T end +# The reconstruction term alone, with `x̂` compared with `target` and normalised by `x`. +function _charbonnier_rec_eval( + loss::CharbonnierRelativeLoss, x̂, target, x, ::Type{T}) where {T} + eps_T = T(loss.ε) + eps_sq = eps_T * eps_T + num_rec = sqrt(_frob_sq_diff(x̂, target) + eps_sq) - eps_T + (num_rec / sqrt(_frob_sq_norm(x) + eps_sq))::T +end + Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, x̂, x) T = promote_type(_state_eltype(x̂), _state_eltype(x), _leaf_eltype(ps.core)) @@ -113,10 +146,8 @@ end Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, x) T = promote_type(_state_eltype(x), _leaf_eltype(ps.core)) - z = encode(model, ps, modality, x) - x̂ = decode(model, ps, modality, z) - z_hat = encode(model, ps, modality, x̂) - _charbonnier_eval(loss, x̂, x, z_hat, z)::T + x̂ = decode(model, ps, modality, encode(model, ps, modality, x)) + _charbonnier_rec_eval(loss, x̂, _reconstruction_target(model, modality, x), x, T) end Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( @@ -160,8 +191,12 @@ Multi-step latent rollout, decoded trajectory reconstruction, and latent Hamilto \right), ``` -where ``z_k = \mathrm{encode}(nn, m, x_k)``, ``\hat{z}_1 = z_1``, ``\hat{z}_{k+1} = \mathrm{integrator}(\hat{z}_k)`` for ``k = 1, \dots, K - 1``, and `H` is a callable latent Hamiltonian (default `nothing`, which uses the canonical quadratic energy ``\frac{1}{2}\|\hat{z}\|_2^2``). +where ``z_k = \mathrm{encode}(nn, m, x_k)``, ``\hat{z}_1 = z_1``, ``\hat{z}_{k+1} = \mathrm{integrator}(\hat{z}_k)`` for ``k = 1, \dots, K - 1``, and `H` is a callable latent Hamiltonian. +`H(ẑ)` takes the latent state as `encode` and `integrator` give it, a vector of length ``2n`` or a matrix of size ``(2n, B)`` for a batched snapshot, and returns one scalar, one value for a whole batch. +The decoded state is compared with ``x_k`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality is recombined into the full field. The weights default to `α_lat = 1.0`, `α_rec = 1.0`, and `α_H = 0.1`. + +`H` has no default: the loss requires a latent Hamiltonian `H` when `α_H > 0`, and the constructor raises an `ArgumentError` without one. A fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, `H` can be `nothing`, and the loss then has no energy term. The data, the output of `integrator` and the value of `H` have the element type of the network's parameters; another element type raises an `ArgumentError`. """ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss @@ -179,14 +214,29 @@ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss H = nothing) (α_lat >= 0 && α_rec >= 0 && α_H >= 0) || throw(ArgumentError( "IsometricRolloutLoss requires non-negative weights α_lat, α_rec, α_H, got α_lat = $α_lat, α_rec = $α_rec, α_H = $α_H.")) + (iszero(α_H) || H !== nothing) || throw(ArgumentError( + "IsometricRolloutLoss requires a latent Hamiltonian H when α_H > 0, got α_H = $α_H and H = nothing. Pass H, or set α_H = 0.")) T = promote_type(typeof(float(α_lat)), typeof(float(α_rec)), typeof(float(α_H))) new{typeof(integrator), typeof(H), T}( integrator, H, T(α_lat), T(α_rec), T(α_H)) end end -_eval_latent_H(::Nothing, z::AbstractArray{T}) where {T} = T(0.5) * sum(abs2, z) -_eval_latent_H(H_fn, z::AbstractArray) = H_fn(z) +# The value of `H` at the first latent state, and the penalty `|H(ẑ) - H(ẑ₁)|²` of a later one. +# `H = nothing` occurs only with `α_H = 0`, and then `H` is never called. +_initial_latent_H(::Nothing, z, ::Type{T}) where {T} = nothing +function _initial_latent_H(H, z, ::Type{T}) where {T} + H1 = H(z) + _check_eltype("IsometricRolloutLoss", "the value of H", H1, T) + H1 +end + +_latent_H_penalty(::Nothing, z, ::Nothing, ::Type{T}) where {T} = zero(T) +function _latent_H_penalty(H, z, H1, ::Type{T}) where {T} + Hk = H(z) + _check_eltype("IsometricRolloutLoss", "the value of H", Hk, T) + abs2(Hk - H1) +end function _check_traj_len(K::Int) K >= 2 || throw(ArgumentError( @@ -219,8 +269,7 @@ Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( T = _leaf_eltype(ps.core) _check_eltype("IsometricRolloutLoss", "the data", xs[1], T) z1 = encode(model, ps, modality, xs[1]) - H1 = _eval_latent_H(loss.H, z1) - _check_eltype("IsometricRolloutLoss", "the value of H", H1, T) + H1 = _initial_latent_H(loss.H, z1, T) z_curr = z1 err_lat = zero(T) err_rec = zero(T) @@ -232,10 +281,9 @@ Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( _check_eltype("IsometricRolloutLoss", "the integrator's output", z_curr, T) err_lat = err_lat + _frob_sq_diff(z_curr, z_k) x_rec_k = decode(model, ps, modality, z_curr) - err_rec = err_rec + _frob_sq_diff(x_rec_k, xs[k]) - Hk = _eval_latent_H(loss.H, z_curr) - _check_eltype("IsometricRolloutLoss", "the value of H", Hk, T) - err_H = err_H + abs2(Hk - H1) + err_rec = err_rec + + _frob_sq_diff(x_rec_k, _reconstruction_target(model, modality, xs[k])) + err_H = err_H + _latent_H_penalty(loss.H, z_curr, H1, T) end inv_steps = inv(T(K - 1)) res = T(loss.α_lat) * err_lat * inv_steps + @@ -273,36 +321,49 @@ Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( end @doc raw""" - CrossModalCycleLoss(mod_a::Symbol, mod_b::Symbol; β_align = 1.0, β_cycle = 0.1) + CrossModalCycleLoss(mod_a::Symbol, mod_b::Symbol; β_align = 1.0, β_cross = 1.0) -Cross-modal paired latent alignment and cycle consistency loss between modalities `mod_a` and `mod_b`: +Cross-modal paired latent alignment and cross-reconstruction loss between modalities `mod_a` and `mod_b`: ```math \mathcal{L}_{\mathrm{cross}}(x_a, x_b) = \beta_{\mathrm{align}}\,\|z_a - z_b\|_F^2 -+ \beta_{\mathrm{cycle}}\,\|\mathrm{encode}(nn, b, \mathrm{decode}(nn, b, z_a)) - z_a\|_F^2, ++ \beta_{\mathrm{cross}}\,\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2, ``` where ``z_a = \mathrm{encode}(nn, a, x_a)`` and ``z_b = \mathrm{encode}(nn, b, x_b)``. -The weights default to `β_align = 1.0` and `β_cycle = 0.1`. +The pair ``(x_a, x_b)`` is the same physical state in two modalities: the alignment term asks that both encode to one latent state, and the cross-reconstruction term asks that the latent state of ``x_a`` decodes to ``x_b``. +The loss is not symmetric in ``a`` and ``b``: the cross-reconstruction term decodes only into modality `mod_b`, so `CrossModalCycleLoss(mod_b, mod_a)` is another loss. +`decode(nn, b, z_a)` is compared with ``x_b`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality is recombined into the full field. +A batch is the sum of its samples' losses; `x_a` and `x_b` with different numbers of samples raise a `DimensionMismatch`. + +The weights are non-negative and default to `β_align = 1.0` and `β_cross = 1.0`; a negative or `NaN` weight raises an `ArgumentError`. `x_a` and `x_b` have the element type of the network's parameters; another element type raises an `ArgumentError`. """ struct CrossModalCycleLoss{T <: Real} <: NetworkLoss mod_a::Symbol mod_b::Symbol β_align::T - β_cycle::T + β_cross::T function CrossModalCycleLoss( mod_a::Symbol, mod_b::Symbol; - β_align::Real = 1.0, β_cycle::Real = 0.1) - (β_align >= 0 && β_cycle >= 0) || throw(ArgumentError( - "CrossModalCycleLoss requires non-negative weights β_align and β_cycle, got β_align = $β_align, β_cycle = $β_cycle.")) - T = promote_type(typeof(float(β_align)), typeof(float(β_cycle))) - new{T}(mod_a, mod_b, T(β_align), T(β_cycle)) + β_align::Real = 1.0, β_cross::Real = 1.0) + (β_align >= 0 && β_cross >= 0) || throw(ArgumentError( + "CrossModalCycleLoss requires non-negative weights β_align and β_cross, got β_align = $β_align, β_cross = $β_cross.")) + T = promote_type(typeof(float(β_align)), typeof(float(β_cross))) + new{T}(mod_a, mod_b, T(β_align), T(β_cross)) end end +function _check_sample_count(z_a, z_b) + size(z_a) == size(z_b) || throw(DimensionMismatch( + "CrossModalCycleLoss: x_a and x_b hold different numbers of samples: their latent states have sizes $(size(z_a)) and $(size(z_b)).")) + nothing +end + +ChainRulesCore.@non_differentiable _check_sample_count(::Any, ::Any) + Base.@constprop :aggressive function _cross_modal_eval( loss::CrossModalCycleLoss, model::MultiModalSymplecticLayer, ps::NamedTuple, x_a, x_b) @@ -311,10 +372,11 @@ Base.@constprop :aggressive function _cross_modal_eval( _check_eltype("CrossModalCycleLoss", "x_b", x_b, T) z_a = encode(model, ps, loss.mod_a, x_a) z_b = encode(model, ps, loss.mod_b, x_b) - x_b_rec = decode(model, ps, loss.mod_b, z_a) - z_a_cyc = encode(model, ps, loss.mod_b, x_b_rec) + _check_sample_count(z_a, z_b) + x̂_b = decode(model, ps, loss.mod_b, z_a) + target = _reconstruction_target(model, loss.mod_b, x_b) res = T(loss.β_align) * _frob_sq_diff(z_a, z_b) + - T(loss.β_cycle) * _frob_sq_diff(z_a_cyc, z_a) + T(loss.β_cross) * _frob_sq_diff(x̂_b, target) res::T end diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 8235a81bb..fddc20e54 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -289,16 +289,26 @@ end @test l_charb_def.λ_lat == 1e-2 step_id = z -> z - l_iso_def = IsometricRolloutLoss(step_id) + H_lat_def = z -> sum(abs2, z) / 2 + l_iso_def = IsometricRolloutLoss(step_id; H = H_lat_def) @test l_iso_def.α_lat == 1.0 @test l_iso_def.α_rec == 1.0 @test l_iso_def.α_H == 0.1 + @test l_iso_def.H === H_lat_def + + # There is no default latent Hamiltonian: α_H > 0 requires H, and α_H = 0 needs none. + re_iso_H = r"^(?:ArgumentError: )?IsometricRolloutLoss requires a latent Hamiltonian H when α_H > 0, got α_H = 0\.1 and H = nothing\. Pass H, or set α_H = 0\.$" + @test_throws re_iso_H IsometricRolloutLoss(step_id) + @test_throws re_iso_H IsometricRolloutLoss(step_id; α_H = 0.1, H = nothing) + @test IsometricRolloutLoss(step_id; α_H = 0.0).H === nothing l_cyc_def = CrossModalCycleLoss(:particles, :grid) @test l_cyc_def.mod_a === :particles @test l_cyc_def.mod_b === :grid @test l_cyc_def.β_align == 1.0 - @test l_cyc_def.β_cycle == 0.1 + @test l_cyc_def.β_cross == 1.0 + @test !hasfield(typeof(l_cyc_def), :β_cycle) + @test_throws MethodError CrossModalCycleLoss(:particles, :grid; β_cycle = 0.1) # Edge validation for loss hyperparameters re_charb_eps = r"^(?:ArgumentError: )?CharbonnierRelativeLoss requires ε > 0 and λ_lat >= 0, got ε = 0\.0, λ_lat = 0\.01\.$" @@ -309,8 +319,12 @@ end re_iso = r"^(?:ArgumentError: )?IsometricRolloutLoss requires non-negative weights α_lat, α_rec, α_H, got α_lat = -1\.0, α_rec = 1\.0, α_H = 0\.1\.$" @test_throws re_iso IsometricRolloutLoss(step_id; α_lat = -1.0) - re_cyc = r"^(?:ArgumentError: )?CrossModalCycleLoss requires non-negative weights β_align and β_cycle, got β_align = 1\.0, β_cycle = -0\.1\.$" - @test_throws re_cyc CrossModalCycleLoss(:particles, :grid; β_cycle = -0.1) + re_cross = r"^(?:ArgumentError: )?CrossModalCycleLoss requires non-negative weights β_align and β_cross, got β_align = 1\.0, β_cross = -0\.1\.$" + @test_throws re_cross CrossModalCycleLoss(:particles, :grid; β_cross = -0.1) + re_cross_nan = r"^(?:ArgumentError: )?CrossModalCycleLoss requires non-negative weights β_align and β_cross, got β_align = 1\.0, β_cross = NaN\.$" + @test_throws re_cross_nan CrossModalCycleLoss(:particles, :grid; β_cross = NaN) + re_align = r"^(?:ArgumentError: )?CrossModalCycleLoss requires non-negative weights β_align and β_cross, got β_align = -0\.1, β_cross = 1\.0\.$" + @test_throws re_align CrossModalCycleLoss(:particles, :grid; β_align = -0.1) @testset "Forward evaluation, @inferred, zero edge case, and exact-zero loss ($T)" for T in REAL_ELTYPES rng = Random.MersenneTwister(202) @@ -333,16 +347,26 @@ end z_hat = encode(nn, :grid, F_hat) eps_T = T(1e-2) lam_T = T(5e-2) - expected_c = (sqrt(sum(abs2, F_hat .- F) + eps_T^2) - eps_T) / - sqrt(sum(abs2, F) + eps_T^2) + + expected_c_rec = (sqrt(sum(abs2, F_hat .- F) + eps_T^2) - eps_T) / + sqrt(sum(abs2, F) + eps_T^2) + expected_c = expected_c_rec + lam_T * sum(abs2, z_hat .- z) / (sum(abs2, F) + eps_T^2) + # The self-reconstruction form has no latent term; the form with an external x̂ has it. val_c1 = @inferred loss_c(nn, :grid, F) val_c2 = @inferred loss_c(nn, :grid, F_hat, F) @test eltype(val_c1) === T && eltype(val_c2) === T - @test isapprox(val_c1, expected_c; atol = 10 * eps(T), rtol = 10 * eps(T)) + @test isapprox(val_c1, expected_c_rec; atol = 10 * eps(T), rtol = 10 * eps(T)) @test isapprox(val_c2, expected_c; atol = 10 * eps(T), rtol = 10 * eps(T)) + # In the self-reconstruction form, a pre-split field (q = F⁺, p = F⁻) gives the value of + # the full field F: the decoded field is compared with F⁺ and F⁻ recombined. + F_plus = (F[:, 1:(nv ÷ 2)] .+ F[:, nv:-1:(nv ÷ 2 + 1)]) ./ sqrt(T(2)) + F_minus = (F[:, 1:(nv ÷ 2)] .- F[:, nv:-1:(nv ÷ 2 + 1)]) ./ sqrt(T(2)) + F_split = (q = F_plus, p = F_minus) + @test isapprox(@inferred(loss_c(nn, :grid, F_split)), val_c1; + atol = 10 * eps(T), rtol = 10 * eps(T)) + # Independent perturbation F_pert != F_hat so z_hat != z and the latent penalty term is active (catches a dropped latent penalty) F_pert = F .+ T(0.5) .* ones(T, nx, nv) z_pert = encode(nn, :grid, F_pert) @@ -391,9 +415,27 @@ end expected_iso = T(1.0) * err_lat + T(0.5) * err_rec + T(0.2) * err_H @test isapprox(val_iso, expected_iso; atol = 15 * eps(T), rtol = 15 * eps(T)) - # CrossModalCycleLoss forward & @inferred + # With α_H = 0 and no H, the loss is the α_lat and α_rec terms alone. + loss_iso_noH = IsometricRolloutLoss(step_lin; α_lat = T(1.0), α_rec = T(0.5), α_H = 0) + val_iso_noH = @inferred loss_iso_noH(nn, :grid, F_traj) + @test eltype(val_iso_noH) === T + @test isapprox(val_iso_noH, T(1.0) * err_lat + T(0.5) * err_rec; + atol = 15 * eps(T), rtol = 15 * eps(T)) + + # A trajectory of pre-split fields gives the value of the trajectory of full fields. + split_traj = [( + q = (F_k[:, 1:(nv ÷ 2)] .+ F_k[:, nv:-1:(nv ÷ 2 + 1)]) ./ + sqrt(T(2)), + p = (F_k[:, 1:(nv ÷ 2)] .- F_k[:, nv:-1:(nv ÷ 2 + 1)]) ./ + sqrt(T(2))) + for F_k in F_traj] + @test isapprox(@inferred(loss_iso(nn, :grid, split_traj)), val_iso; + atol = 15 * eps(T), rtol = 15 * eps(T)) + + # CrossModalCycleLoss forward & @inferred: the alignment term and the + # cross-reconstruction term ‖decode(b, z_a) - x_b‖². loss_cyc = CrossModalCycleLoss( - :particles, :grid; β_align = T(1.2), β_cycle = T(0.3) + :particles, :grid; β_align = T(1.2), β_cross = T(0.3) ) x_p = (q = randn(rng, T, 1, 3), p = randn(rng, T, 1, 3)) val_cyc = @inferred loss_cyc(nn, x_p, F) @@ -401,21 +443,69 @@ end z_a = encode(nn, :particles, x_p) z_b = encode(nn, :grid, F) - z_a_cyc = encode(nn, :grid, decode(nn, :grid, z_a)) - expected_cyc = T(1.2) * sum(abs2, z_a .- z_b) + T(0.3) * sum(abs2, z_a_cyc .- z_a) + expected_cyc = T(1.2) * sum(abs2, z_a .- z_b) + + T(0.3) * sum(abs2, decode(nn, :grid, z_a) .- F) @test isapprox(val_cyc, expected_cyc; atol = 10 * eps(T), rtol = 10 * eps(T)) - # Reverse direction CrossModalCycleLoss(:grid, :particles) where :particles is only chamber-wise right-invertible (checks the cycle term against the explicit encode-decode round trip) + # Reverse direction CrossModalCycleLoss(:grid, :particles): the particle decoder returns + # the flat state [vec(q); vec(p)]. loss_cyc_rev = CrossModalCycleLoss( - :grid, :particles; β_align = T(1.2), β_cycle = T(0.5)) + :grid, :particles; β_align = T(1.2), β_cross = T(0.5)) val_cyc_rev = @inferred loss_cyc_rev(nn, F, x_p) - z_b_cyc = encode(nn, :particles, decode(nn, :particles, z_b)) expected_cyc_rev = T(1.2) * sum(abs2, z_b .- z_a) + - T(0.5) * sum(abs2, z_b_cyc .- z_b) + T(0.5) * sum( + abs2, decode(nn, :particles, z_b) .- vcat(vec(x_p.q), vec(x_p.p))) @test isapprox( - val_cyc_rev, expected_cyc_rev; atol = 10 * eps(T), rtol = 10 * - eps(T) - ) + val_cyc_rev, expected_cyc_rev; atol = 10 * eps(T), rtol = 10 * eps(T)) + + # The three receivers, each with x_a and x_b and with the pair (x_a, x_b). + ps_nt, ps_np = params(nn), nn.params + @test (@inferred loss_cyc(nn, (x_p, F))) == val_cyc + @test (@inferred loss_cyc(nn.model, ps_nt, x_p, F)) == val_cyc + @test (@inferred loss_cyc(nn.model, ps_nt, (x_p, F))) == val_cyc + @test (@inferred loss_cyc(nn.model, ps_np, x_p, F)) == val_cyc + @test (@inferred loss_cyc(nn.model, ps_np, (x_p, F))) == val_cyc + + # A batch is the sum over its samples, with no division by B. + B = 3 + x_p_b = (q = randn(rng, T, 1, 3, B), p = randn(rng, T, 1, 3, B)) + F_b = randn(rng, T, nx, nv, B) + val_cyc_b = @inferred loss_cyc(nn, x_p_b, F_b) + @test eltype(val_cyc_b) === T + val_cyc_sum = sum( + b -> loss_cyc(nn, (q = x_p_b.q[:, :, b], p = x_p_b.p[:, :, b]), F_b[:, :, b]), + 1:B) + @test isapprox(val_cyc_b, val_cyc_sum; rtol = 10 * eps(T)) + @test (@inferred loss_cyc(nn, (x_p_b, F_b))) == val_cyc_b + + # x_a and x_b with different numbers of samples raise a DimensionMismatch. + re_samples = r"^(?:DimensionMismatch: )?CrossModalCycleLoss: x_a and x_b hold different numbers of samples: their latent states have sizes \(2,\) and \(2, 3\)\.$" + @test_throws re_samples loss_cyc(nn, x_p, F_b) + re_samples_rev = r"^(?:DimensionMismatch: )?CrossModalCycleLoss: x_a and x_b hold different numbers of samples: their latent states have sizes \(2, 3\) and \(2,\)\.$" + @test_throws re_samples_rev loss_cyc(nn, x_p_b, F) + + # A pre-split field (q = F⁺, p = F⁻) for x_b gives the value of the full field F. The + # parity components come from explicit parity matrices: vec(F⁺) = P₊ vec(F). + half = nv ÷ 2 + P_plus = zeros(T, nx * half, nx * nv) + P_minus = zeros(T, nx * half, nx * nv) + for j in 1:half, i in 1:nx + + row = (j - 1) * nx + i + col_top = (j - 1) * nx + i + col_bot = (nv - j) * nx + i + P_plus[row, col_top] = P_plus[row, col_bot] = inv(sqrt(T(2))) + P_minus[row, col_top] = inv(sqrt(T(2))) + P_minus[row, col_bot] = -inv(sqrt(T(2))) + end + split_of(F2) = (q = reshape(P_plus * vec(F2), nx, half), + p = reshape(P_minus * vec(F2), nx, half)) + @test isapprox(@inferred(loss_cyc(nn, x_p, split_of(F))), val_cyc; + rtol = 10 * eps(T)) + F_b_split = (q = cat((split_of(F_b[:, :, b]).q for b in 1:B)...; dims = 3), + p = cat((split_of(F_b[:, :, b]).p for b in 1:B)...; dims = 3)) + @test isapprox(@inferred(loss_cyc(nn, x_p_b, F_b_split)), val_cyc_b; + rtol = 10 * eps(T)) # Another precision than the network's raises an ArgumentError that names both. S = T === Float32 ? Float64 : Float32 @@ -473,7 +563,7 @@ end H_fn = ξ -> 0.5 * sum(abs2, ξ) + 0.05 * sum(ξ .^ 4) loss_iso = IsometricRolloutLoss( step_fn; α_lat = 1.0, α_rec = 0.8, α_H = 0.25, H = H_fn) - loss_cyc = CrossModalCycleLoss(:particles, :grid; β_align = 1.0, β_cycle = 0.2) + loss_cyc = CrossModalCycleLoss(:particles, :grid; β_align = 1.0, β_cross = 0.2) flat0, rebuild_ps = _flatten_params(nn.params) h = 1e-6 @@ -663,3 +753,138 @@ end @test all(isfinite, _grad_vec(g_sym)) end end + +# The untied core of the control networks below. +function _untied_core() + SymplecticAutoencoder( + 4, 2; n_encoder_blocks = 1, n_decoder_blocks = 1, n_encoder_layers = 1, + n_decoder_layers = 1, n_decoder_output_layers = 1, sympnet_upscale = 2) +end + +@testset "Cross-reconstruction term against a dense linear oracle ($T)" for T in REAL_ELTYPES + rng = Random.MersenneTwister(505) + Nx, Nv, mx, mv = 4, 4, 2, 1 + half = Nv ÷ 2 + s = inv(sqrt(T(2))) + # Exactly orthonormal leaves, with entries 0, ±1 and ±1/√2. + Φx = T[s 0; -s 0; 0 1; 0 0] + Φv = reshape(T[s, s], half, mv) + Φc = reshape(T[s, -s], 2, 1) + + fe = BilinearStiefelLayer(Nx, Nv, mx, mv) + core = WeightTiedSymplecticAutoencoder( + 4, 2; n_blocks = 1, n_layers = 0, n_ambient_layers = 0, sympnet_upscale = 2) + nn = NeuralNetwork( + MultiModalSymplecticAutoencoder((a = fe, b = fe), core), CPU(), T; rng = rng) + nn_sa = NeuralNetwork( + MultiModalSymplecticAutoencoder((a = fe, b = fe), _untied_core()), CPU(), T; + rng = rng) + for net in (nn, nn_sa), m in (:a, :b) + + copyto!(parent(getfield(params(net).frontends, m).Φx), Φx) + copyto!(parent(getfield(params(net).frontends, m).Φv), Φv) + end + @test length(params(nn).core) == 1 + copyto!(parent(params(nn).core.L1.weight), Φc) + + # Dense encoder and decoder from the leaves and explicit parity matrices, vec(F) x-fastest: + # vec(F⁺) = P₊ vec(F), vec(F⁻) = P₋ vec(F), vec(Φxᵀ G Φv) = kron(Φv, Φx)ᵀ vec(G), and the + # weight-tied core with no SympNet layers is (q, p) ↦ (Φcᵀ q, Φcᵀ p). + P_plus = zeros(T, Nx * half, Nx * Nv) + P_minus = zeros(T, Nx * half, Nx * Nv) + for j in 1:half, i in 1:Nx + + row = (j - 1) * Nx + i + col_top = (j - 1) * Nx + i + col_bot = (Nv - j) * Nx + i + P_plus[row, col_top] = P_plus[row, col_bot] = s + P_minus[row, col_top] = s + P_minus[row, col_bot] = -s + end + K = kron(Φv, Φx) + O = zeros(T, 1, 2) + C_enc = [transpose(Φc) O; O transpose(Φc)] + Enc = C_enc * vcat(transpose(K) * P_plus, transpose(K) * P_minus) + Dec = hcat(transpose(P_plus) * K, transpose(P_minus) * K) * transpose(C_enc) + + x_a = T(3.7) .* randn(rng, T, Nx, Nv) + x_b = T(3.7) .* randn(rng, T, Nx, Nv) + z_a = Enc * vec(x_a) + z_b = Enc * vec(x_b) + @test isapprox(encode(nn, :a, x_a), z_a; rtol = 10 * eps(T)) + @test isapprox(vec(decode(nn, :b, z_a)), Dec * z_a; rtol = 10 * eps(T)) + + β_align, β_cross = T(0.7), T(1.3) + loss = CrossModalCycleLoss(:a, :b; β_align = β_align, β_cross = β_cross) + expected = β_align * sum(abs2, z_a - z_b) + β_cross * sum(abs2, Dec * z_a - vec(x_b)) + val = @inferred loss(nn, x_a, x_b) + @test val isa T + @test isapprox(val, expected; rtol = 10 * eps(T)) + + # With β_align = 0 the loss is the cross term alone, which is not small and depends on x_b. + loss_cross = CrossModalCycleLoss(:a, :b; β_align = 0, β_cross = 1) + cross = sum(abs2, Dec * z_a - vec(x_b)) + @test isapprox(loss_cross(nn, x_a, x_b), cross; rtol = 10 * eps(T)) + @test cross > 1e-2 + x_b2 = x_b .+ T(0.5) + cross2 = sum(abs2, Dec * z_a - vec(x_b2)) + @test isapprox(loss_cross(nn, x_a, x_b2), cross2; rtol = 10 * eps(T)) + @test abs(loss_cross(nn, x_a, x_b2) - loss_cross(nn, x_a, x_b)) > T(1e-2) * cross + + # The cycle term ‖encode(b, decode(b, z_a)) - z_a‖ is zero for the weight-tied core with a + # right-invertible frontend, and not zero for the untied core. + z_a_nn = encode(nn, :a, x_a) + cycle = norm(encode(nn, :b, decode(nn, :b, z_a_nn)) - z_a_nn) + @test cycle <= 10 * eps(T) * max(one(T), norm(z_a_nn)) + z_a_sa = encode(nn_sa, :a, x_a) + cycle_sa = norm(encode(nn_sa, :b, decode(nn_sa, :b, z_a_sa)) - z_a_sa) + @test cycle_sa > T(1e-3) * norm(z_a_sa) + + # The self-reconstruction form of CharbonnierRelativeLoss has no latent term: on the untied + # core, where that term is not zero, the value is the reconstruction term alone. The form + # with an external x̂ = decode(encode(x)) keeps the latent term. + ε = T(1e-2) + loss_c = CharbonnierRelativeLoss(; ε = ε, λ_lat = 1) + x̂_a = decode(nn_sa, :a, encode(nn_sa, :a, x_a)) + rec = (sqrt(sum(abs2, x̂_a - x_a) + ε^2) - ε) / sqrt(sum(abs2, x_a) + ε^2) + lat = sum(abs2, encode(nn_sa, :a, x̂_a) - z_a_sa) / (sum(abs2, x_a) + ε^2) + @test lat > T(1e-3) * rec + val_self = @inferred loss_c(nn_sa, :a, x_a) + @test val_self isa T + @test isapprox(val_self, rec; rtol = 10 * eps(T)) + @test isapprox(loss_c(nn_sa, :a, x̂_a, x_a), rec + lat; rtol = 10 * eps(T)) +end + +@testset "Docstrings state the limits of the frontends and the losses" begin + # The source text of every docstring of X, with its line breaks as spaces. `Docs.doc` has + # no method outside an interactive session, so the text comes from the module's metadata. + function docstring(X) + binding = Base.Docs.Binding(GeometricMachineLearning, nameof(X)) + multidoc = Base.Docs.meta(GeometricMachineLearning)[binding] + replace(join((join(d.text) for d in values(multidoc.docs)), "\n"), '\n' => ' ') + end + documented = (CanonicalSortLayer, CanonicalSlotLayer, InverseCanonicalSlotLayer, + SymmetricSymplecticAutoencoder, BilinearStiefelLayer, LowRankStiefelLayer, + MultiModalSymplecticAutoencoder, CharbonnierRelativeLoss, IsometricRolloutLoss, + CrossModalCycleLoss) + absent = ("collective canonical slots", "relative modes canonically", + "canonical quadratic energy", "β_cycle", "symplectic or Poisson frontend") + for X in documented + str = docstring(X) + for phrase in absent + @test !occursin(phrase, str) + end + @test !occursin(r"(Proposition|Theorem|Corollary|Section) \d", str) + end + for X in (CanonicalSortLayer, CanonicalSlotLayer, SymmetricSymplecticAutoencoder) + @test occursin("keep their particle order", docstring(X)) + end + @test occursin("fixed particle count", docstring(CanonicalSlotLayer)) + for X in (BilinearStiefelLayer, LowRankStiefelLayer, MultiModalSymplecticAutoencoder) + @test occursin("declare on the parity coordinates", docstring(X)) + end + @test occursin("requires a latent Hamiltonian", docstring(IsometricRolloutLoss)) + @test occursin("same physical state", docstring(CrossModalCycleLoss)) + @test occursin("not symmetric", docstring(CrossModalCycleLoss)) + @test occursin("skips", docstring(CharbonnierRelativeLoss)) +end diff --git a/test/reduced_system/reference_shift_and_vjp.jl b/test/reduced_system/reference_shift_and_vjp.jl index 5bc386f57..130d9c4e7 100644 --- a/test/reduced_system/reference_shift_and_vjp.jl +++ b/test/reduced_system/reference_shift_and_vjp.jl @@ -506,3 +506,70 @@ end v_red0(v_out0, zero(T), q̃, p̃, nothing) @test norm(v_out0 .- v_expected) > 1e-3 * norm(v_expected) end + +@testset "HRedSys with a reference that is not an equilibrium, against the unshifted system in shifted coordinates ($T, $mode)" for T in REAL_ELTYPES, + mode in (:jacobian, :vjp) + + rng = Random.MersenneTwister(808) + N = 6 + n = 2 + arch = WeightTiedSymplecticAutoencoder( + 2 * N, 2 * n; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) + nn = NeuralNetwork(arch, CPU(), T; rng = rng) + + # H(q, p) = ½‖p‖² + ½‖q‖² + ¼ Σ qᵢ⁴, centred at the origin. + v_field!(v, t, q, p, params) = (v .= p; nothing) + f_field!(f, t, q, p, params) = (f .= .-q .- q .^ 3; nothing) + h_field(t, q, p, params) = T(0.5) * (sum(abs2, p) + sum(abs2, q)) + + T(0.25) * sum(x -> x^4, q) + + q_ref = randn(rng, T, N) + p_ref = randn(rng, T, N) + z_ref = vcat(q_ref, p_ref) + # The full vector field (∂H/∂p, -∂H/∂q) does not vanish at the reference. + @test norm(vcat(p_ref, q_ref .+ q_ref .^ 3)) >= 0.1 + q0 = q_ref .+ T(0.3) .* randn(rng, T, N) + p0 = p_ref .+ T(0.3) .* randn(rng, T, N) + + rs = HRedSys( + 2 * N, 2 * n, encoder(nn), decoder(nn), v_field!, f_field!, h_field, + (zero(T), one(T)), T(0.05), (q = StateVariable(q0), p = StateVariable(p0)); + integrator = ImplicitMidpoint(), q_ref = q_ref, p_ref = p_ref, + pullback_mode = mode) + + # The oracle: the system without a reference, with the fields translated by the reference + # and the initial state written in shifted coordinates. + v̂!(v, t, y_q, y_p, params) = v_field!(v, t, y_q .+ q_ref, y_p .+ p_ref, params) + f̂!(f, t, y_q, y_p, params) = f_field!(f, t, y_q .+ q_ref, y_p .+ p_ref, params) + ĥ(t, y_q, y_p, params) = h_field(t, y_q .+ q_ref, y_p .+ p_ref, params) + rs_oracle = HRedSys( + 2 * N, 2 * n, encoder(nn), decoder(nn), v̂!, f̂!, ĥ, + (zero(T), one(T)), T(0.05), + (q = StateVariable(q0 .- q_ref), p = StateVariable(p0 .- p_ref)); + integrator = ImplicitMidpoint(), pullback_mode = mode) + + sol = integrate_reduced_system(rs) + sol_oracle = integrate_reduced_system(rs_oracle) + @test length(sol.q) == 21 # initial state + 20 steps + for k in 0:20 + q̃_k, p̃_k = Vector(sol.q[k]), Vector(sol.p[k]) + ξ_k = vcat(q̃_k, p̃_k) + ξ̂_k = vcat(Vector(sol_oracle.q[k]), Vector(sol_oracle.p[k])) + bound = 50 * eps(T) * max(one(T), norm(ξ̂_k)) + @test norm(ξ_k - ξ̂_k) <= bound + @test norm(decode(rs, ξ_k) - (decoder(nn)(ξ̂_k) .+ z_ref)) <= bound + @test abs(rs.h_reduced(zero(T), q̃_k, p̃_k, nothing) - + rs_oracle.h_reduced(zero(T), q̃_k, p̃_k, nothing)) <= bound + end + # The trajectory moves. + @test norm(vcat(Vector(sol.q[20]), Vector(sol.p[20])) - + vcat(Vector(sol.q[0]), Vector(sol.p[0]))) > T(1e-2) + + # The check can fail: at ξ_0 the shifted reduced field differs from the unshifted one. + q̃_0, p̃_0 = Vector(sol.q[0]), Vector(sol.p[0]) + v_shifted, v_unshifted = zeros(T, n), zeros(T, n) + rs.v_reduced(v_shifted, zero(T), q̃_0, p̃_0, nothing) + build_v_reduced(v_field!, f_field!, decoder(nn); pullback_mode = mode)( + v_unshifted, zero(T), q̃_0, p̃_0, nothing) + @test norm(v_shifted - v_unshifted) > 1e-3 * norm(v_shifted) +end From 44bf1e77b6c8a9db505d3ffb2b35b229e4a359b7 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sat, 31 Oct 2026 00:31:57 +0100 Subject: [PATCH 2/6] Qualify the invariance claim of the particle page and correct two text lines The page's SymmetricSymplecticAutoencoder paragraph states the same distinct-coordinates qualifier as its opening paragraph. The CHANGELOG bullet names only the docstrings that state the frontends' limits, and a test comment says exactly when Docs.doc has no method. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- docs/src/layers/symmetric_particle_layers.md | 2 +- test/architectures/multimodal_symplectic_autoencoder.jl | 3 ++- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d59f74f5..edd5f64bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -154,7 +154,7 @@ breaking release). ### Changed -- **The multi-modal losses of #370 lose the terms that are zero by construction, `IsometricRolloutLoss` has no default latent Hamiltonian, and the docstrings of the particle and parity frontends state their limits.** No release has these names, so nothing is deprecated. `CrossModalCycleLoss` replaces its cycle term ``\|\mathrm{encode}(nn, b, \mathrm{decode}(nn, b, z_a)) - z_a\|_F^2``, which is zero for a weight-tied core with a right-invertible frontend, with the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``, which depends on the data `x_b`; the keyword `β_cycle = 0.1` becomes `β_cross = 1.0`, a passed `β_cycle` raises Julia's `MethodError`, and a negative or `NaN` `β_cross` raises an `ArgumentError`. `x_a` and `x_b` whose latent states have different sizes, such as one sample against a batch, raise a `DimensionMismatch`. The self-reconstruction forms of `CharbonnierRelativeLoss`, those without an external `x̂`, skip the latent term, which is zero there for a weight-tied core with a right-invertible frontend, so `λ_lat` acts only in the form with `x̂`. `IsometricRolloutLoss(integrator)` raises an `ArgumentError`, because `α_H > 0` requires a latent Hamiltonian `H` and the fixed ``\frac{1}{2}\|\hat{z}\|_2^2`` of #370 is not the system's energy; with `α_H = 0` it needs none and calls none. The cross-reconstruction term, the self-reconstruction form of `CharbonnierRelativeLoss` and the `α_rec` term of `IsometricRolloutLoss` compare a decoded `BilinearStiefelLayer` field with a pre-split pair `(q = F⁺, p = F⁻)` after recombining the pair into the full field. The docstrings of `CanonicalSortLayer`, `CanonicalSlotLayer`, `InverseCanonicalSlotLayer`, `SymmetricSymplecticAutoencoder`, `BilinearStiefelLayer`, `LowRankStiefelLayer`, `MultiModalSymplecticAutoencoder` and the three losses, and the pages `docs/src/layers/symmetric_particle_layers.md` and `docs/src/layers/stiefel_layer.md`, state that sorting is discontinuous where two particles cross and suits systems that keep their particle order, that the slot layer is a linear projection of sorted states for a fixed particle count, and that the `(q, p)` of the parity frontends is canonical only for the structure that they declare on the parity coordinates, not for the Vlasov bracket. A new testset of `test/reduced_system/reference_shift_and_vjp.jl` checks `HRedSys` with a reference that is not an equilibrium against the unshifted system in shifted coordinates. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. +- **The multi-modal losses of #370 lose the terms that are zero by construction, `IsometricRolloutLoss` has no default latent Hamiltonian, and the docstrings of the particle and parity frontends state their limits.** No release has these names, so nothing is deprecated. `CrossModalCycleLoss` replaces its cycle term ``\|\mathrm{encode}(nn, b, \mathrm{decode}(nn, b, z_a)) - z_a\|_F^2``, which is zero for a weight-tied core with a right-invertible frontend, with the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``, which depends on the data `x_b`; the keyword `β_cycle = 0.1` becomes `β_cross = 1.0`, a passed `β_cycle` raises Julia's `MethodError`, and a negative or `NaN` `β_cross` raises an `ArgumentError`. `x_a` and `x_b` whose latent states have different sizes, such as one sample against a batch, raise a `DimensionMismatch`. The self-reconstruction forms of `CharbonnierRelativeLoss`, those without an external `x̂`, skip the latent term, which is zero there for a weight-tied core with a right-invertible frontend, so `λ_lat` acts only in the form with `x̂`. `IsometricRolloutLoss(integrator)` raises an `ArgumentError`, because `α_H > 0` requires a latent Hamiltonian `H` and the fixed ``\frac{1}{2}\|\hat{z}\|_2^2`` of #370 is not the system's energy; with `α_H = 0` it needs none and calls none. The cross-reconstruction term, the self-reconstruction form of `CharbonnierRelativeLoss` and the `α_rec` term of `IsometricRolloutLoss` compare a decoded `BilinearStiefelLayer` field with a pre-split pair `(q = F⁺, p = F⁻)` after recombining the pair into the full field. The docstrings of `CanonicalSortLayer`, `CanonicalSlotLayer`, `SymmetricSymplecticAutoencoder`, `BilinearStiefelLayer`, `LowRankStiefelLayer` and `MultiModalSymplecticAutoencoder`, and the pages `docs/src/layers/symmetric_particle_layers.md` and `docs/src/layers/stiefel_layer.md`, state that sorting is discontinuous where two particles cross and suits systems that keep their particle order, that the slot layer is a linear projection of sorted states for a fixed particle count, and that the `(q, p)` of the parity frontends is canonical only for the structure that they declare on the parity coordinates, not for the Vlasov bracket. A new testset of `test/reduced_system/reference_shift_and_vjp.jl` checks `HRedSys` with a reference that is not an equilibrium against the unshifted system in shifted coordinates. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. - **`SymplecticAttentionQ` and `SymplecticAttentionP` accept a single sample.** A `(q, p)` `NamedTuple` of vectors, and a plain vector, are reshaped to one column, passed through the matrix method and returned as vectors, so the forward layer accepts the same inputs as its inverse. On `main` a `NamedTuple` of vectors and a plain vector raised `MethodError: no method matching _custom_transpose(::Vector{Float64})`. The test is `test/layers/symplectic_attention.jl`. diff --git a/docs/src/layers/symmetric_particle_layers.md b/docs/src/layers/symmetric_particle_layers.md index eeb87887f..dc6790f62 100644 --- a/docs/src/layers/symmetric_particle_layers.md +++ b/docs/src/layers/symmetric_particle_layers.md @@ -45,7 +45,7 @@ InverseCanonicalSlotLayer ## Symmetric Symplectic Autoencoder -[`SymmetricSymplecticAutoencoder`](@ref) combines a chosen particle frontend (`:sort`, `:pairwise`, or `:slot`) with [`WeightTiedSymplecticAutoencoder`](@ref). Only `:sort` and `:slot` give a permutation-invariant latent state; `:pairwise` is equivariant only in its frontend, because the core is not equivariant. +[`SymmetricSymplecticAutoencoder`](@ref) combines a chosen particle frontend (`:sort`, `:pairwise`, or `:slot`) with [`WeightTiedSymplecticAutoencoder`](@ref). Only `:sort` and `:slot` give a latent state that is invariant under a relabelling of particles with distinct coordinates along the sorting axis; `:pairwise` is equivariant only in its frontend, because the core is not equivariant. ```@docs SymmetricSymplecticAutoencoder diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index fddc20e54..662aae4f7 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -857,7 +857,8 @@ end @testset "Docstrings state the limits of the frontends and the losses" begin # The source text of every docstring of X, with its line breaks as spaces. `Docs.doc` has - # no method outside an interactive session, so the text comes from the module's metadata. + # no method until the REPL standard library loads, so the text comes from the module's + # metadata. function docstring(X) binding = Base.Docs.Binding(GeometricMachineLearning, nameof(X)) multidoc = Base.Docs.meta(GeometricMachineLearning)[binding] From 17b645097c1a1717be20d9b4ff095553598de2c2 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sat, 31 Oct 2026 08:45:26 +0100 Subject: [PATCH 3/6] Fix the review findings of the multi-modal losses and frontend docstrings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - CharbonnierRelativeLoss: the form with an external x̂ recombines a pre-split BilinearStiefelLayer pair, given as x or as x̂, as the self form does. A full decoded field against a pre-split x gave a silently different value. The docstring says that the self form skips the latent term for every core and frontend, and that the term need not be zero for CanonicalSortLayer and CanonicalSlotLayer. - IsometricRolloutLoss: with α_H = 0 the constructor stores nothing in place of a given H, so H is never called and an Inf value of H no longer makes the loss NaN. - CrossModalCycleLoss: the DimensionMismatch names different latent sizes, not different sample counts, which was wrong for (2,) and (2, 1). - BilinearStiefelLayer, LowRankStiefelLayer, stiefel_layer.md: the linearised Vlasov bracket gives a {Q, P} that is not the identity and can be singular; it is singular only for some bases. - CHANGELOG: the Changed bullet folds into the Added bullets, without the false claim that the dropped terms are zero by construction. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 8 +++----- docs/src/layers/stiefel_layer.md | 2 +- src/layers/bilinear_stiefel_layer.jl | 3 ++- src/layers/low_rank_stiefel_layer.jl | 3 ++- src/loss/multimodal_losses.jl | 18 ++++++++++-------- .../multimodal_symplectic_autoencoder.jl | 10 ++++++++-- 6 files changed, 26 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index edd5f64bb..22345e6ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -146,16 +146,14 @@ breaking release). - **`scripts/utilities/metal_abort_runs.jl` counts whether one of the two Metal process aborts recurs, one fresh process per counted run, and `scripts/utilities/metal_abort_session.jl` runs the arm that needs two requests to one session.** Two processes died in the first Adam call of a structured network on `MetalBackend()` in `Float32` — `LASympNet(4)` from a command-encoder assertion while a second evaluation read a Julia vector, and `VolumePreservingFeedForward(4)` from host-heap corruption detected in the LLVM JIT — and neither had a reproducer when the runs began. **One death of the session arm has the signature of the `LASympNet(4)` item:** run 18 died with `SIGSEGV` in `AGX::ComputeContext…::endComputePass` under `-[AGXG16XFamilyCommandBuffer commit]` and Metal.jl's `commit_with_queue_key!`, from `flush_batched_queues!` inside a task's `device_synchronize`, with no GeometricMachineLearning frame in the stack. It happened in the load phase, before the overlap stamps and outside the arm's counted window, so it is recorded as a death and not as a reproduction, and the arm stays open. A counted run prints its pid, `MTL_DEBUG_LAYER`, `Threads.nthreads()`, the version and path of every package of the stack, `Metal.functional()` and the device's name before any Metal work; it takes a forward pass and a Zygote pullback on one network and times the *first* training call of a new one with `@timed Base.invokelatest`, so that the compilation is inside the timer, and a `compile_time` of 0 makes the run void. `N` is 59 per arm, and a run's crash report counts only when the report's `"pid"` is the run's own, so the sixteen reports of other Julia processes in `~/Library/Logs/DiagnosticReports` are not counted; `audit` recounts an arm from its logs and the reports directory. Neither arm that a shell process runs aborted in 59 of 59 runs, on Julia 1.13.3 and Metal.jl 1.11.1, the Metal.jl of the aborts, and the third arm, a private session's two `ex` calls, has 16 counted runs, all clean, and the one recorded death. The session script appends `M8-TRAIN-START` and `M8-READ-START` to the run's log file before the timed call can finish, `audit` counts a death of that arm as its abort only when they show the overlap, and a child that dies before it prints its pid is counted from the pid that the parent logs. An issue for `JuliaGPU/Metal.jl` is drafted and not filed, on the hypothesis that `flush_batched_queues!` commits another task's open command buffer with no lock. `KNOWN_ISSUES.md` K35 records the counts, both stacks, the procedure and the draft. -- **`CanonicalSortLayer`, `InvertCanonicalSortLayer`, `PairwiseSymplecticLayer`, `PairwiseSymplecticLayerQ`, `PairwiseSymplecticLayerP`, `InversePairwiseSymplecticLayerQ`, `InversePairwiseSymplecticLayerP`, `CanonicalSlotLayer`, `InverseCanonicalSlotLayer`, and `SymmetricSymplecticAutoencoder` add particle symplectic layers that are invariant or equivariant under relabelling of the particles, and a weight-tied symplectic autoencoder with three particle frontends.** All ten names are exported. `CanonicalSortLayer(M, d, axis = 1)` makes `CanonicalSortLayer{N, M, Axis}` with `N = 2Md`, sorts `M` particles in ``\mathbb{R}^d`` along coordinate row `axis` of `q` using a stable merge sort (`MergeSort`), applies the same permutation ``\sigma \in S_M`` to both `q` and `p`, and returns `(q = q_sorted, p = p_sorted, perm = σ)` on `(q, p)` `NamedTuple` inputs of shape `(d, M)` or `(d, M, B)` (or the sorted state on flattened `2Md` vectors and `(2Md, B)` matrices). `InvertCanonicalSortLayer(M, d)` makes `InvertCanonicalSortLayer{N, M}` and applies ``\sigma^{-1}`` (`invperm`) to restore the original particle ordering; `inverse_layer` pairs it with `CanonicalSortLayer`. `CanonicalSortLayer`, `InvertCanonicalSortLayer`, `CanonicalSlotLayer` and `InverseCanonicalSlotLayer` carry the phase-space dimension of the flattened state as their first type parameter, `N = 2Md`, so `input_dimension` and `output_dimension` return the length of that state, and their inputs are checked against the fixed `d`, with the first dimension of an input exactly its input dimension. `PairwiseSymplecticLayerQ(M, d, upscale, σ)` and `PairwiseSymplecticLayerP(M, d, upscale, σ)` apply ``S_M``-equivariant canonical shears combining a self term ``K_{\mathrm{self}}^T \sigma(K_{\mathrm{self}} u_i + b_{\mathrm{self}})`` and a skew-symmetric pairwise relative potential term ``\frac{1}{M}\sum_{j=1}^M K_{\mathrm{pair}}^T \tau(K_{\mathrm{pair}}(u_i - u_j))`` with ``\tau(x) = \sigma(x) - \sigma(-x)``, conserving total canonical momentum (or position) when ``K_{\mathrm{self}} = 0`` and pairing with `InversePairwiseSymplecticLayerQ` and `InversePairwiseSymplecticLayerP` through `inverse_layer`. `CanonicalSlotLayer(M, d, NSlots)` (`M >= 2`, `1 <= NSlots <= M`, `CanonicalSlotLayer{N, NOut, M, NSlots}` with `NOut = 2 NSlots d`) first sorts the particles along coordinate row 1 of `q` with `CanonicalSortLayer`, so its slots are invariant under a relabelling of particles with distinct first coordinates. It then separates the center-of-mass mode ``(Q_{\mathrm{cm}}, P_{\mathrm{cm}}) = \frac{1}{\sqrt{M}}\sum_{i=1}^M (q_i, p_i)`` from the ``M - 1`` zero-sum Helmert relative modes, sorts the ``M - 1`` relative modes canonically by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2`` (`MergeSort`), and projects the sorted relative modes onto ``N_{\mathrm{slots}} - 1`` slots via a learned `StiefelManifold` parameter ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)`` (`NSlots == 1` keeps only the center-of-mass slot with no parameter leaf). `InverseCanonicalSlotLayer(M, d, NSlots)` takes the same type parameters, with input dimension `NOut` and output dimension `N`, and lifts the slots back symplectically via ``\Phi_{\mathrm{rel}}`` and the transpose Helmert matrix. `SymmetricSymplecticAutoencoder(M, d, reduced_dim; frontend = :sort | :pairwise | :slot, ...) <: SymplecticCompression` combines one of the three particle frontends with the weight-tied `WeightTiedSymplecticAutoencoder` core. With `:sort` and `:slot` the latent state is invariant under a relabelling of particles with distinct coordinates along the sorting axis (`sort_axis` for `:sort`, coordinate 1 for `:slot`). With `:pairwise` only the frontend is permutation-equivariant, because the core is not, so the latent state is not invariant. Right-inversion ``\mathrm{encoder}(nn)(\mathrm{decoder}(nn)(\xi)) = \xi`` holds to machine precision for every ``\xi`` with `:pairwise`; with `:sort` and `:slot` it holds only for ``\xi`` whose decoded state lies in the canonical chamber, where the sorting permutations on the decoded state are the identity, because the decoder does not map into that chamber in general. The tests are `test/layers/symmetric_particle_layers.jl`. +- **`CanonicalSortLayer`, `InvertCanonicalSortLayer`, `PairwiseSymplecticLayer`, `PairwiseSymplecticLayerQ`, `PairwiseSymplecticLayerP`, `InversePairwiseSymplecticLayerQ`, `InversePairwiseSymplecticLayerP`, `CanonicalSlotLayer`, `InverseCanonicalSlotLayer`, and `SymmetricSymplecticAutoencoder` add particle symplectic layers that are invariant or equivariant under relabelling of the particles, and a weight-tied symplectic autoencoder with three particle frontends.** All ten names are exported. `CanonicalSortLayer(M, d, axis = 1)` makes `CanonicalSortLayer{N, M, Axis}` with `N = 2Md`, sorts `M` particles in ``\mathbb{R}^d`` along coordinate row `axis` of `q` using a stable merge sort (`MergeSort`), applies the same permutation ``\sigma \in S_M`` to both `q` and `p`, and returns `(q = q_sorted, p = p_sorted, perm = σ)` on `(q, p)` `NamedTuple` inputs of shape `(d, M)` or `(d, M, B)` (or the sorted state on flattened `2Md` vectors and `(2Md, B)` matrices). `InvertCanonicalSortLayer(M, d)` makes `InvertCanonicalSortLayer{N, M}` and applies ``\sigma^{-1}`` (`invperm`) to restore the original particle ordering; `inverse_layer` pairs it with `CanonicalSortLayer`. `CanonicalSortLayer`, `InvertCanonicalSortLayer`, `CanonicalSlotLayer` and `InverseCanonicalSlotLayer` carry the phase-space dimension of the flattened state as their first type parameter, `N = 2Md`, so `input_dimension` and `output_dimension` return the length of that state, and their inputs are checked against the fixed `d`, with the first dimension of an input exactly its input dimension. `PairwiseSymplecticLayerQ(M, d, upscale, σ)` and `PairwiseSymplecticLayerP(M, d, upscale, σ)` apply ``S_M``-equivariant canonical shears combining a self term ``K_{\mathrm{self}}^T \sigma(K_{\mathrm{self}} u_i + b_{\mathrm{self}})`` and a skew-symmetric pairwise relative potential term ``\frac{1}{M}\sum_{j=1}^M K_{\mathrm{pair}}^T \tau(K_{\mathrm{pair}}(u_i - u_j))`` with ``\tau(x) = \sigma(x) - \sigma(-x)``, conserving total canonical momentum (or position) when ``K_{\mathrm{self}} = 0`` and pairing with `InversePairwiseSymplecticLayerQ` and `InversePairwiseSymplecticLayerP` through `inverse_layer`. `CanonicalSlotLayer(M, d, NSlots)` (`M >= 2`, `1 <= NSlots <= M`, `CanonicalSlotLayer{N, NOut, M, NSlots}` with `NOut = 2 NSlots d`) first sorts the particles along coordinate row 1 of `q` with `CanonicalSortLayer`, so its slots are invariant under a relabelling of particles with distinct first coordinates. It then separates the center-of-mass mode ``(Q_{\mathrm{cm}}, P_{\mathrm{cm}}) = \frac{1}{\sqrt{M}}\sum_{i=1}^M (q_i, p_i)`` from the ``M - 1`` zero-sum Helmert relative modes, sorts the ``M - 1`` relative modes canonically by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2`` (`MergeSort`), and projects the sorted relative modes onto ``N_{\mathrm{slots}} - 1`` slots via a learned `StiefelManifold` parameter ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)`` (`NSlots == 1` keeps only the center-of-mass slot with no parameter leaf). `InverseCanonicalSlotLayer(M, d, NSlots)` takes the same type parameters, with input dimension `NOut` and output dimension `N`, and lifts the slots back symplectically via ``\Phi_{\mathrm{rel}}`` and the transpose Helmert matrix. `SymmetricSymplecticAutoencoder(M, d, reduced_dim; frontend = :sort | :pairwise | :slot, ...) <: SymplecticCompression` combines one of the three particle frontends with the weight-tied `WeightTiedSymplecticAutoencoder` core. With `:sort` and `:slot` the latent state is invariant under a relabelling of particles with distinct coordinates along the sorting axis (`sort_axis` for `:sort`, coordinate 1 for `:slot`). With `:pairwise` only the frontend is permutation-equivariant, because the core is not, so the latent state is not invariant. Right-inversion ``\mathrm{encoder}(nn)(\mathrm{decoder}(nn)(\xi)) = \xi`` holds to machine precision for every ``\xi`` with `:pairwise`; with `:sort` and `:slot` it holds only for ``\xi`` whose decoded state lies in the canonical chamber, where the sorting permutations on the decoded state are the identity, because the decoder does not map into that chamber in general. Sorting is discontinuous where two particles cross, so the `:sort` and `:slot` frontends suit systems that keep their particle order, and the slot layer is a linear projection of the sorted states for a fixed particle count; the docstrings state both limits. The tests are `test/layers/symmetric_particle_layers.jl`. -- **`BilinearStiefelLayer`, `TransposeBilinearStiefelLayer`, `LowRankStiefelLayer`, `DecodeLowRankStiefelLayer`, `FastDLRATraceLoss`, and `solve!` on `BilinearStiefelLayer` add two-sided velocity-parity Stiefel frontends, a gauge-invariant dynamical low-rank Stiefel frontend and decoder, a fast low-rank trace loss, and Tucker-2 HOSVD initialization.** The five new types are exported. `BilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)` (`Nx >= 1`, even `Nv >= 2`, `1 <= mx <= Nx`, `1 <= mv <= Nv ÷ 2`) allocates `(Φx = StiefelManifold(...), Φv = StiefelManifold(...))` with ``\Phi_x \in \mathrm{St}(m_x, N_x)`` and ``\Phi_v \in \mathrm{St}(m_v, N_v / 2)``, splits an input array ``F`` of shape `(Nx, Nv)` or `(Nx, Nv, B)` into velocity-even ``F^+ = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} + F_{:, N_v:-1:N_v/2+1})`` and velocity-odd ``F^- = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} - F_{:, N_v:-1:N_v/2+1})`` components (or accepts a pre-split `(q, p)` `NamedTuple` of shape `(Nx, Nv ÷ 2)` or `(Nx, Nv ÷ 2, B)`), applies optional Cholesky mass factors `(Rx, Rv)` when `metric_factors` is supplied, and projects ``Q_{\mathrm{red}} = \Phi_x^T F^+ \Phi_v`` and ``P_{\mathrm{red}} = \Phi_x^T F^- \Phi_v`` to return vectorized `(q = vec(Q_red), p = vec(P_red))` of half-dimension ``n = m_x m_v``; `TransposeBilinearStiefelLayer` (`inverse_layer(l)`) is its exact right-inverse. `LowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)` takes a DLRA `NamedTuple` `(X = X, S = S, V = V)` of arbitrary rank ``r_{\mathrm{in}} \ge 1`` and evaluates ``Q_{\mathrm{red}} = (\Phi_x^T X) S ((V^+)^T \Phi_v)`` and ``P_{\mathrm{red}} = (\Phi_x^T X) S ((V^-)^T \Phi_v)`` in ``O(r_{\mathrm{in}}(m_x N_x + m_v N_v))`` operations without forming ``X S V^T``, matching `BilinearStiefelLayer` on ``F = X S V^T`` identically and remaining invariant under ``O(r_{\mathrm{in}}) \times O(r_{\mathrm{in}})`` gauge transformations, while `DecodeLowRankStiefelLayer` (`inverse_layer(l)`) computes the truncated SVD of ``[Q_{\mathrm{red}}\;\; P_{\mathrm{red}}]`` of rank ``r_{\mathrm{eff}} = \min(r_{\mathrm{out}}, m_x, 2 m_v)`` and QR re-orthonormalizes the lifted factors to return `(X = X_hat, S = S_hat, V = V_hat)` on the Stiefel manifold. `FastDLRATraceLoss()` evaluates ``\|X S V^T - \hat{X} \hat{S} \hat{V}^T\|_F^2`` in ``O(r^2(N_x + N_v))`` operations via Gram-matrix traces clamped at `zero(T)`. `solve!(nn, F_snapshots)` and `solve!(l, ps, F_snapshots)` initialize `Φx` and `Φv` on a `BilinearStiefelLayer` from the top-``m_x`` spatial and top-``m_v`` parity-velocity left singular vectors of `F_snapshots`. The tests are `test/layers/bilinear_and_low_rank_stiefel.jl`. +- **`BilinearStiefelLayer`, `TransposeBilinearStiefelLayer`, `LowRankStiefelLayer`, `DecodeLowRankStiefelLayer`, `FastDLRATraceLoss`, and `solve!` on `BilinearStiefelLayer` add two-sided velocity-parity Stiefel frontends, a gauge-invariant dynamical low-rank Stiefel frontend and decoder, a fast low-rank trace loss, and Tucker-2 HOSVD initialization.** The five new types are exported. `BilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)` (`Nx >= 1`, even `Nv >= 2`, `1 <= mx <= Nx`, `1 <= mv <= Nv ÷ 2`) allocates `(Φx = StiefelManifold(...), Φv = StiefelManifold(...))` with ``\Phi_x \in \mathrm{St}(m_x, N_x)`` and ``\Phi_v \in \mathrm{St}(m_v, N_v / 2)``, splits an input array ``F`` of shape `(Nx, Nv)` or `(Nx, Nv, B)` into velocity-even ``F^+ = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} + F_{:, N_v:-1:N_v/2+1})`` and velocity-odd ``F^- = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} - F_{:, N_v:-1:N_v/2+1})`` components (or accepts a pre-split `(q, p)` `NamedTuple` of shape `(Nx, Nv ÷ 2)` or `(Nx, Nv ÷ 2, B)`), applies optional Cholesky mass factors `(Rx, Rv)` when `metric_factors` is supplied, and projects ``Q_{\mathrm{red}} = \Phi_x^T F^+ \Phi_v`` and ``P_{\mathrm{red}} = \Phi_x^T F^- \Phi_v`` to return vectorized `(q = vec(Q_red), p = vec(P_red))` of half-dimension ``n = m_x m_v``; `TransposeBilinearStiefelLayer` (`inverse_layer(l)`) is its exact right-inverse. `LowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)` takes a DLRA `NamedTuple` `(X = X, S = S, V = V)` of arbitrary rank ``r_{\mathrm{in}} \ge 1`` and evaluates ``Q_{\mathrm{red}} = (\Phi_x^T X) S ((V^+)^T \Phi_v)`` and ``P_{\mathrm{red}} = (\Phi_x^T X) S ((V^-)^T \Phi_v)`` in ``O(r_{\mathrm{in}}(m_x N_x + m_v N_v))`` operations without forming ``X S V^T``, matching `BilinearStiefelLayer` on ``F = X S V^T`` identically and remaining invariant under ``O(r_{\mathrm{in}}) \times O(r_{\mathrm{in}})`` gauge transformations, while `DecodeLowRankStiefelLayer` (`inverse_layer(l)`) computes the truncated SVD of ``[Q_{\mathrm{red}}\;\; P_{\mathrm{red}}]`` of rank ``r_{\mathrm{eff}} = \min(r_{\mathrm{out}}, m_x, 2 m_v)`` and QR re-orthonormalizes the lifted factors to return `(X = X_hat, S = S_hat, V = V_hat)` on the Stiefel manifold. `FastDLRATraceLoss()` evaluates ``\|X S V^T - \hat{X} \hat{S} \hat{V}^T\|_F^2`` in ``O(r^2(N_x + N_v))`` operations via Gram-matrix traces clamped at `zero(T)`. `solve!(nn, F_snapshots)` and `solve!(l, ps, F_snapshots)` initialize `Φx` and `Φv` on a `BilinearStiefelLayer` from the top-``m_x`` spatial and top-``m_v`` parity-velocity left singular vectors of `F_snapshots`. The `(q, p)` of the parity frontends is canonical only for the structure that they declare on the parity coordinates, not for the Vlasov bracket, and the docstrings say so. The tests are `test/layers/bilinear_and_low_rank_stiefel.jl`. -- **`MultiModalSymplecticAutoencoder`, `encode`, `decode`, `translate`, `CharbonnierRelativeLoss`, `IsometricRolloutLoss`, and `CrossModalCycleLoss` add multi-modal symplectic compression and multi-modal training losses, and `HRedSys` gains reference-state shifting, reverse-mode VJP vector fields, and preallocated work buffers.** All seven new names are exported. `MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression` combines a non-empty `NamedTuple` of modality frontends (individual layers or `Chain`s of symplectic particle layers) whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` (with `modality` a `Symbol` or `Val(modality)`) applies `frontends[modality]` followed by `encoder(core)`, `decode(nn, modality, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])`, and `translate(nn, from, to, x)` computes `decode(nn, to, encode(nn, from, x))`, raising an `ArgumentError` on an empty or dimension-mismatched `frontends` or an unknown modality symbol. `CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)` evaluates, for an external reconstruction `x̂`, the exact-zero Charbonnier relative reconstruction loss ``(\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon) / \sqrt{\|x\|_F^2 + \varepsilon^2} + \lambda_{\mathrm{lat}} \|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2 / (\|x\|_F^2 + \varepsilon^2)``, and without `x̂` the reconstruction term alone for ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` (supporting low-rank `(X, S, V)` inputs via `FastDLRATraceLoss`), `IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.1, H = nothing)` evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty, with an `ArgumentError` when `α_H > 0` and `H = nothing`, and `CrossModalCycleLoss(mod_a, mod_b; β_align = 1.0, β_cross = 1.0)` evaluates paired latent alignment ``\|z_a - z_b\|_F^2`` plus the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``. `IsometricRolloutLoss` raises an `ArgumentError` naming both element types when the data, the output of `integrator` or `H` has another element type than the network's parameters, and `CrossModalCycleLoss` does so for `x_a` or `x_b`. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (converted to the state element type, shifting the decoder lift by ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` and subtracting ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` before `rs.encoder` in `integrate_reduced_system`) and `pullback_mode::Symbol = :jacobian | :vjp`, where `:vjp` evaluates ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` in one reverse-mode pullback `Zygote.pullback(decoder, ξ)` and both modes use the preallocated buffers of `cache` (`ξ_buf`, `v_intermediate`, `f_intermediate`, and `jac_result` for `:jacobian`, with a per-element-type `IdDict{DataType, Any}` of caches for `ForwardDiff.Dual` Newton states). The default `cache` of `build_v_reduced` and `build_f_reduced` has the element type of the decoder's parameters, and the default `cache` of `HRedSys` has the element type of `ics.q`; a state of another floating-point element type than the cache's raises an `ArgumentError`. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. +- **`MultiModalSymplecticAutoencoder`, `encode`, `decode`, `translate`, `CharbonnierRelativeLoss`, `IsometricRolloutLoss`, and `CrossModalCycleLoss` add multi-modal symplectic compression and multi-modal training losses, and `HRedSys` gains reference-state shifting, reverse-mode VJP vector fields, and preallocated work buffers.** All seven new names are exported. `MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression` combines a non-empty `NamedTuple` of modality frontends (individual layers or `Chain`s of symplectic particle layers) whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` (with `modality` a `Symbol` or `Val(modality)`) applies `frontends[modality]` followed by `encoder(core)`, `decode(nn, modality, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])`, and `translate(nn, from, to, x)` computes `decode(nn, to, encode(nn, from, x))`, raising an `ArgumentError` on an empty or dimension-mismatched `frontends` or an unknown modality symbol. `CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)` evaluates, for an external reconstruction `x̂`, the exact-zero Charbonnier relative reconstruction loss ``(\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon) / \sqrt{\|x\|_F^2 + \varepsilon^2} + \lambda_{\mathrm{lat}} \|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2 / (\|x\|_F^2 + \varepsilon^2)``, and without `x̂` the reconstruction term alone for ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` (supporting low-rank `(X, S, V)` inputs via `FastDLRATraceLoss`), where `λ_lat` has no effect. To get the latent term for a modality, pass ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` to the form with an external `x̂`. The latent term is zero for a weight-tied core with a right-invertible frontend, and need not be zero otherwise, including for a `CanonicalSortLayer` frontend and for a `CanonicalSlotLayer` frontend with more than one slot. Both forms of `CharbonnierRelativeLoss`, the cross-reconstruction term of `CrossModalCycleLoss` and the `α_rec` term of `IsometricRolloutLoss` compare fields in the form that `decode` returns, so a pre-split pair `(q = F⁺, p = F⁻)` of a `BilinearStiefelLayer` modality is recombined into the full field first. `IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.1, H = nothing)` evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty. It has no default latent Hamiltonian: `α_H > 0` requires `H`, and the constructor raises an `ArgumentError` without one, because a fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, `H` can be `nothing`, and the loss never calls it. `CrossModalCycleLoss(mod_a, mod_b; β_align = 1.0, β_cross = 1.0)` evaluates paired latent alignment ``\|z_a - z_b\|_F^2`` plus the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``. It has no cycle term, because ``\|\mathrm{encode}(nn, b, \mathrm{decode}(nn, b, z_a)) - z_a\|_F^2`` is zero for a weight-tied core with a right-invertible frontend and not in general. A negative or `NaN` weight raises an `ArgumentError`, and `x_a` and `x_b` whose latent states have different sizes, such as one sample against a batch, raise a `DimensionMismatch`. `IsometricRolloutLoss` raises an `ArgumentError` naming both element types when the data, the output of `integrator` or `H` has another element type than the network's parameters, and `CrossModalCycleLoss` does so for `x_a` or `x_b`. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (converted to the state element type, shifting the decoder lift by ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` and subtracting ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` before `rs.encoder` in `integrate_reduced_system`) and `pullback_mode::Symbol = :jacobian | :vjp`, where `:vjp` evaluates ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` in one reverse-mode pullback `Zygote.pullback(decoder, ξ)` and both modes use the preallocated buffers of `cache` (`ξ_buf`, `v_intermediate`, `f_intermediate`, and `jac_result` for `:jacobian`, with a per-element-type `IdDict{DataType, Any}` of caches for `ForwardDiff.Dual` Newton states). The default `cache` of `build_v_reduced` and `build_f_reduced` has the element type of the decoder's parameters, and the default `cache` of `HRedSys` has the element type of `ics.q`; a state of another floating-point element type than the cache's raises an `ArgumentError`. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. ### Changed -- **The multi-modal losses of #370 lose the terms that are zero by construction, `IsometricRolloutLoss` has no default latent Hamiltonian, and the docstrings of the particle and parity frontends state their limits.** No release has these names, so nothing is deprecated. `CrossModalCycleLoss` replaces its cycle term ``\|\mathrm{encode}(nn, b, \mathrm{decode}(nn, b, z_a)) - z_a\|_F^2``, which is zero for a weight-tied core with a right-invertible frontend, with the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``, which depends on the data `x_b`; the keyword `β_cycle = 0.1` becomes `β_cross = 1.0`, a passed `β_cycle` raises Julia's `MethodError`, and a negative or `NaN` `β_cross` raises an `ArgumentError`. `x_a` and `x_b` whose latent states have different sizes, such as one sample against a batch, raise a `DimensionMismatch`. The self-reconstruction forms of `CharbonnierRelativeLoss`, those without an external `x̂`, skip the latent term, which is zero there for a weight-tied core with a right-invertible frontend, so `λ_lat` acts only in the form with `x̂`. `IsometricRolloutLoss(integrator)` raises an `ArgumentError`, because `α_H > 0` requires a latent Hamiltonian `H` and the fixed ``\frac{1}{2}\|\hat{z}\|_2^2`` of #370 is not the system's energy; with `α_H = 0` it needs none and calls none. The cross-reconstruction term, the self-reconstruction form of `CharbonnierRelativeLoss` and the `α_rec` term of `IsometricRolloutLoss` compare a decoded `BilinearStiefelLayer` field with a pre-split pair `(q = F⁺, p = F⁻)` after recombining the pair into the full field. The docstrings of `CanonicalSortLayer`, `CanonicalSlotLayer`, `SymmetricSymplecticAutoencoder`, `BilinearStiefelLayer`, `LowRankStiefelLayer` and `MultiModalSymplecticAutoencoder`, and the pages `docs/src/layers/symmetric_particle_layers.md` and `docs/src/layers/stiefel_layer.md`, state that sorting is discontinuous where two particles cross and suits systems that keep their particle order, that the slot layer is a linear projection of sorted states for a fixed particle count, and that the `(q, p)` of the parity frontends is canonical only for the structure that they declare on the parity coordinates, not for the Vlasov bracket. A new testset of `test/reduced_system/reference_shift_and_vjp.jl` checks `HRedSys` with a reference that is not an equilibrium against the unshifted system in shifted coordinates. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. - - **`SymplecticAttentionQ` and `SymplecticAttentionP` accept a single sample.** A `(q, p)` `NamedTuple` of vectors, and a plain vector, are reshaped to one column, passed through the matrix method and returned as vectors, so the forward layer accepts the same inputs as its inverse. On `main` a `NamedTuple` of vectors and a plain vector raised `MethodError: no method matching _custom_transpose(::Vector{Float64})`. The test is `test/layers/symplectic_attention.jl`. - **`SymplecticAutoencoder` builds with one encoder or decoder block, and names the limit when there are too many.** `compute_iterations_for_symplectic_system`, which `SymplecticAutoencoder` and `WeightTiedSymplecticAutoencoder` share, divided by `n_blocks - 1` and raised `DivideError` for `n_blocks = 1`; it now returns `[reduced_dim, full_dim]`, one block. A block count whose spacing `(full_dim - reduced_dim) ÷ 2 ÷ (n_blocks - 1)` rounds to zero raised Base's `ArgumentError: step cannot be zero`, and now raises an `ArgumentError` that names the largest allowed count: `SymplecticAutoencoder(8, 4; n_decoder_blocks = 4)` gives "n_blocks (4) must be at most 3 for full_dim (8) and reduced_dim (4).", and a block count below 1 gives "n_blocks (0) must be at least 1.". As before, the `SymplecticAutoencoder` constructor accepts these block counts, and the error is raised when the network is built. The tests are `test/architectures/symplectic_autoencoder_tests.jl` and `test/architectures/weight_tied_symplectic_autoencoder.jl`. diff --git a/docs/src/layers/stiefel_layer.md b/docs/src/layers/stiefel_layer.md index 8ea9cbf71..8d70e2cee 100644 --- a/docs/src/layers/stiefel_layer.md +++ b/docs/src/layers/stiefel_layer.md @@ -8,7 +8,7 @@ dynamical low-rank approximation (DLRA) states ``(X, S, V)``. Their ``(q, p)`` is canonical only for the structure ``\mathbb{J}`` that they declare on the parity coordinates. A frontend that is linear in ``f`` is not a Poisson map for the Vlasov bracket: with ``L^2``-orthonormal even and odd bases, the bracket linearised at ``f_0`` gives a -singular ``\{Q, P\}``, not the identity. +``\{Q, P\}`` that is not the identity and can be singular. ## Bilinear Grid and Spline Stiefel Layers diff --git a/src/layers/bilinear_stiefel_layer.jl b/src/layers/bilinear_stiefel_layer.jl index d60bb5dbb..ef7964750 100644 --- a/src/layers/bilinear_stiefel_layer.jl +++ b/src/layers/bilinear_stiefel_layer.jl @@ -33,7 +33,8 @@ mapped into the Euclidean metric by `(Rx, Rv)` prior to Stiefel projection. The coordinates ``(q, p)`` are canonical only for the structure ``\mathbb{J}_{2n}`` that the layer and [`LowRankStiefelLayer`](@ref) declare on the parity coordinates. A frontend that is linear in ``F`` is not a Poisson map for the Vlasov bracket: with ``L^2``-orthonormal even and -odd bases, the bracket linearised at ``f_0`` gives a singular ``\{Q, P\}``, not the identity. +odd bases, the bracket linearised at ``f_0`` gives a ``\{Q, P\}`` that is not the identity +and can be singular. See also [`TransposeBilinearStiefelLayer`](@ref), [`LowRankStiefelLayer`](@ref), and [`inverse_layer`](@ref). diff --git a/src/layers/low_rank_stiefel_layer.jl b/src/layers/low_rank_stiefel_layer.jl index 4638252da..d55be9904 100644 --- a/src/layers/low_rank_stiefel_layer.jl +++ b/src/layers/low_rank_stiefel_layer.jl @@ -27,7 +27,8 @@ orthogonal gauge transformations ``(X, S, V) \mapsto (X O_x, O_x^T S O_v, V O_v) The coordinates ``(q, p)`` are canonical only for the structure ``\mathbb{J}_{2n}`` that the layer and [`BilinearStiefelLayer`](@ref) declare on the parity coordinates. A frontend that is linear in ``f = X S V^T`` is not a Poisson map for the Vlasov bracket: with ``L^2``-orthonormal even and -odd bases, the bracket linearised at ``f_0`` gives a singular ``\{Q, P\}``, not the identity. +odd bases, the bracket linearised at ``f_0`` gives a ``\{Q, P\}`` that is not the identity +and can be singular. See also [`DecodeLowRankStiefelLayer`](@ref), [`BilinearStiefelLayer`](@ref), [`FastDLRATraceLoss`](@ref), and [`inverse_layer`](@ref). diff --git a/src/loss/multimodal_losses.jl b/src/loss/multimodal_losses.jl index c9d8288ce..f017ac6ef 100644 --- a/src/loss/multimodal_losses.jl +++ b/src/loss/multimodal_losses.jl @@ -67,7 +67,7 @@ end Smooth Charbonnier relative reconstruction loss, with a latent consistency term when the reconstruction ``\hat{x}`` is passed. -Called as `loss(nn, m, x̂, x)` with an external ``\hat{x}`` of the same form as ``x``, it evaluates +Called as `loss(nn, m, x̂, x)` with an external ``\hat{x}``, it evaluates ```math \mathcal{L}_{\mathrm{Charb}}(\hat{x}, x) = @@ -83,8 +83,8 @@ Called as `loss(nn, m, x)` (or `loss(nn, x)` for the first modality), it evaluat \qquad \hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x)), ``` -and compares ``\hat{x}`` with ``x`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality is recombined into the full field. -In this form the latent term is zero for a weight-tied core with a right-invertible frontend, because then ``\mathrm{encode}(nn, m, \hat{x}) = \mathrm{encode}(nn, m, x)``, so the method skips it and `λ_lat` has no effect. For any other core, pass ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` to the form with an external ``\hat{x}`` to get the latent term. +Both forms compare ``\hat{x}`` with ``x`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality, passed as ``x`` or as ``\hat{x}``, is recombined into the full field. +The form without ``\hat{x}`` skips the latent term, so `λ_lat` has no effect there. The term is zero for a weight-tied core with a right-invertible frontend, because then ``\mathrm{encode}(nn, m, \hat{x}) = \mathrm{encode}(nn, m, x)``. It need not be zero for any other core, nor for a [`CanonicalSortLayer`](@ref) or [`CanonicalSlotLayer`](@ref) frontend, whose decoder does not map into the chamber of sorted states in general. To get the latent term there, pass ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` to the form with an external ``\hat{x}``. Here ``\varepsilon > 0`` (default `1e-3`) regularizes the Frobenius norm at ``\hat{x} = x`` and ``x = 0``, and ``\lambda_{\mathrm{lat}} \ge 0`` (default `1e-2`) weights the latent re-encoding consistency penalty. @@ -130,7 +130,8 @@ Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( T = promote_type(_state_eltype(x̂), _state_eltype(x), _leaf_eltype(ps.core)) z = encode(model, ps, modality, x) z_hat = encode(model, ps, modality, x̂) - _charbonnier_eval(loss, x̂, x, z_hat, z)::T + _charbonnier_eval(loss, _reconstruction_target(model, modality, x̂), + _reconstruction_target(model, modality, x), z_hat, z)::T end Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( @@ -196,7 +197,7 @@ where ``z_k = \mathrm{encode}(nn, m, x_k)``, ``\hat{z}_1 = z_1``, ``\hat{z}_{k+1 The decoded state is compared with ``x_k`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality is recombined into the full field. The weights default to `α_lat = 1.0`, `α_rec = 1.0`, and `α_H = 0.1`. -`H` has no default: the loss requires a latent Hamiltonian `H` when `α_H > 0`, and the constructor raises an `ArgumentError` without one. A fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, `H` can be `nothing`, and the loss then has no energy term. +`H` has no default: the loss requires a latent Hamiltonian `H` when `α_H > 0`, and the constructor raises an `ArgumentError` without one. A fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, `H` can be `nothing`, and the loss then has no energy term: it stores `nothing` in place of a given `H` and never calls it. The data, the output of `integrator` and the value of `H` have the element type of the network's parameters; another element type raises an `ArgumentError`. """ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss @@ -217,8 +218,9 @@ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss (iszero(α_H) || H !== nothing) || throw(ArgumentError( "IsometricRolloutLoss requires a latent Hamiltonian H when α_H > 0, got α_H = $α_H and H = nothing. Pass H, or set α_H = 0.")) T = promote_type(typeof(float(α_lat)), typeof(float(α_rec)), typeof(float(α_H))) - new{typeof(integrator), typeof(H), T}( - integrator, H, T(α_lat), T(α_rec), T(α_H)) + H_used = iszero(α_H) ? nothing : H + new{typeof(integrator), typeof(H_used), T}( + integrator, H_used, T(α_lat), T(α_rec), T(α_H)) end end @@ -358,7 +360,7 @@ end function _check_sample_count(z_a, z_b) size(z_a) == size(z_b) || throw(DimensionMismatch( - "CrossModalCycleLoss: x_a and x_b hold different numbers of samples: their latent states have sizes $(size(z_a)) and $(size(z_b)).")) + "CrossModalCycleLoss: the latent states of x_a and x_b have different sizes, $(size(z_a)) and $(size(z_b)).")) nothing end diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 662aae4f7..910923f1f 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -301,6 +301,8 @@ end @test_throws re_iso_H IsometricRolloutLoss(step_id) @test_throws re_iso_H IsometricRolloutLoss(step_id; α_H = 0.1, H = nothing) @test IsometricRolloutLoss(step_id; α_H = 0.0).H === nothing + # With α_H = 0 a given H is not stored, so it is never called. + @test IsometricRolloutLoss(step_id; α_H = 0.0, H = H_lat_def).H === nothing l_cyc_def = CrossModalCycleLoss(:particles, :grid) @test l_cyc_def.mod_a === :particles @@ -366,6 +368,10 @@ end F_split = (q = F_plus, p = F_minus) @test isapprox(@inferred(loss_c(nn, :grid, F_split)), val_c1; atol = 10 * eps(T), rtol = 10 * eps(T)) + # The form with an external x̂ recombines a pre-split x and a pre-split x̂ in the same way. + @test isapprox(@inferred(loss_c(nn, :grid, F_hat, F_split)), val_c2; + atol = 10 * eps(T), rtol = 10 * eps(T)) + @test @inferred(loss_c(nn, :grid, F_split, F_split)) === zero(T) # Independent perturbation F_pert != F_hat so z_hat != z and the latent penalty term is active (catches a dropped latent penalty) F_pert = F .+ T(0.5) .* ones(T, nx, nv) @@ -479,9 +485,9 @@ end @test (@inferred loss_cyc(nn, (x_p_b, F_b))) == val_cyc_b # x_a and x_b with different numbers of samples raise a DimensionMismatch. - re_samples = r"^(?:DimensionMismatch: )?CrossModalCycleLoss: x_a and x_b hold different numbers of samples: their latent states have sizes \(2,\) and \(2, 3\)\.$" + re_samples = r"^(?:DimensionMismatch: )?CrossModalCycleLoss: the latent states of x_a and x_b have different sizes, \(2,\) and \(2, 3\)\.$" @test_throws re_samples loss_cyc(nn, x_p, F_b) - re_samples_rev = r"^(?:DimensionMismatch: )?CrossModalCycleLoss: x_a and x_b hold different numbers of samples: their latent states have sizes \(2, 3\) and \(2,\)\.$" + re_samples_rev = r"^(?:DimensionMismatch: )?CrossModalCycleLoss: the latent states of x_a and x_b have different sizes, \(2, 3\) and \(2,\)\.$" @test_throws re_samples_rev loss_cyc(nn, x_p_b, F) # A pre-split field (q = F⁺, p = F⁻) for x_b gives the value of the full field F. The From a0d19463c477d483a18ac02dcf9c79fdd3921850 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sat, 31 Oct 2026 09:04:34 +0100 Subject: [PATCH 4/6] =?UTF-8?q?Restore=20the=20Charbonnier=20latent=20term?= =?UTF-8?q?,=20default=20=CE=B1=5FH=20to=200,=20rename=20CrossModalLoss?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - CharbonnierRelativeLoss: the self form evaluates the latent term again, with x̂ = decode(encode(x)), so both forms evaluate one formula and λ_lat acts in both. The skipped term is not zero for CanonicalSortLayer or CanonicalSlotLayer frontends. - A pre-split BilinearStiefelLayer pair whose halves have different element types raises an ArgumentError naming both, in place of a MethodError from _parity_recombine (CodeRabbit comment 4431005164). - IsometricRolloutLoss: α_H defaults to 0.0, so IsometricRolloutLoss(integrator) constructs and has no energy term. - CrossModalCycleLoss is renamed CrossModalLoss; it has no cycle term. No release has either name. - KNOWN_ISSUES K41: the name and its line citations follow. - CHANGELOG: the unreleased bullet describes the end state. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- KNOWN_ISSUES.md | 10 +- docs/src/reduced_order_modeling/losses.md | 2 +- src/GeometricMachineLearning.jl | 2 +- src/loss/multimodal_losses.jl | 78 ++++++------- .../multimodal_symplectic_autoencoder.jl | 106 +++++++++++------- 6 files changed, 110 insertions(+), 90 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 22345e6ad..5b65227d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -150,7 +150,7 @@ breaking release). - **`BilinearStiefelLayer`, `TransposeBilinearStiefelLayer`, `LowRankStiefelLayer`, `DecodeLowRankStiefelLayer`, `FastDLRATraceLoss`, and `solve!` on `BilinearStiefelLayer` add two-sided velocity-parity Stiefel frontends, a gauge-invariant dynamical low-rank Stiefel frontend and decoder, a fast low-rank trace loss, and Tucker-2 HOSVD initialization.** The five new types are exported. `BilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)` (`Nx >= 1`, even `Nv >= 2`, `1 <= mx <= Nx`, `1 <= mv <= Nv ÷ 2`) allocates `(Φx = StiefelManifold(...), Φv = StiefelManifold(...))` with ``\Phi_x \in \mathrm{St}(m_x, N_x)`` and ``\Phi_v \in \mathrm{St}(m_v, N_v / 2)``, splits an input array ``F`` of shape `(Nx, Nv)` or `(Nx, Nv, B)` into velocity-even ``F^+ = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} + F_{:, N_v:-1:N_v/2+1})`` and velocity-odd ``F^- = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} - F_{:, N_v:-1:N_v/2+1})`` components (or accepts a pre-split `(q, p)` `NamedTuple` of shape `(Nx, Nv ÷ 2)` or `(Nx, Nv ÷ 2, B)`), applies optional Cholesky mass factors `(Rx, Rv)` when `metric_factors` is supplied, and projects ``Q_{\mathrm{red}} = \Phi_x^T F^+ \Phi_v`` and ``P_{\mathrm{red}} = \Phi_x^T F^- \Phi_v`` to return vectorized `(q = vec(Q_red), p = vec(P_red))` of half-dimension ``n = m_x m_v``; `TransposeBilinearStiefelLayer` (`inverse_layer(l)`) is its exact right-inverse. `LowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)` takes a DLRA `NamedTuple` `(X = X, S = S, V = V)` of arbitrary rank ``r_{\mathrm{in}} \ge 1`` and evaluates ``Q_{\mathrm{red}} = (\Phi_x^T X) S ((V^+)^T \Phi_v)`` and ``P_{\mathrm{red}} = (\Phi_x^T X) S ((V^-)^T \Phi_v)`` in ``O(r_{\mathrm{in}}(m_x N_x + m_v N_v))`` operations without forming ``X S V^T``, matching `BilinearStiefelLayer` on ``F = X S V^T`` identically and remaining invariant under ``O(r_{\mathrm{in}}) \times O(r_{\mathrm{in}})`` gauge transformations, while `DecodeLowRankStiefelLayer` (`inverse_layer(l)`) computes the truncated SVD of ``[Q_{\mathrm{red}}\;\; P_{\mathrm{red}}]`` of rank ``r_{\mathrm{eff}} = \min(r_{\mathrm{out}}, m_x, 2 m_v)`` and QR re-orthonormalizes the lifted factors to return `(X = X_hat, S = S_hat, V = V_hat)` on the Stiefel manifold. `FastDLRATraceLoss()` evaluates ``\|X S V^T - \hat{X} \hat{S} \hat{V}^T\|_F^2`` in ``O(r^2(N_x + N_v))`` operations via Gram-matrix traces clamped at `zero(T)`. `solve!(nn, F_snapshots)` and `solve!(l, ps, F_snapshots)` initialize `Φx` and `Φv` on a `BilinearStiefelLayer` from the top-``m_x`` spatial and top-``m_v`` parity-velocity left singular vectors of `F_snapshots`. The `(q, p)` of the parity frontends is canonical only for the structure that they declare on the parity coordinates, not for the Vlasov bracket, and the docstrings say so. The tests are `test/layers/bilinear_and_low_rank_stiefel.jl`. -- **`MultiModalSymplecticAutoencoder`, `encode`, `decode`, `translate`, `CharbonnierRelativeLoss`, `IsometricRolloutLoss`, and `CrossModalCycleLoss` add multi-modal symplectic compression and multi-modal training losses, and `HRedSys` gains reference-state shifting, reverse-mode VJP vector fields, and preallocated work buffers.** All seven new names are exported. `MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression` combines a non-empty `NamedTuple` of modality frontends (individual layers or `Chain`s of symplectic particle layers) whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` (with `modality` a `Symbol` or `Val(modality)`) applies `frontends[modality]` followed by `encoder(core)`, `decode(nn, modality, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])`, and `translate(nn, from, to, x)` computes `decode(nn, to, encode(nn, from, x))`, raising an `ArgumentError` on an empty or dimension-mismatched `frontends` or an unknown modality symbol. `CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)` evaluates, for an external reconstruction `x̂`, the exact-zero Charbonnier relative reconstruction loss ``(\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon) / \sqrt{\|x\|_F^2 + \varepsilon^2} + \lambda_{\mathrm{lat}} \|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2 / (\|x\|_F^2 + \varepsilon^2)``, and without `x̂` the reconstruction term alone for ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` (supporting low-rank `(X, S, V)` inputs via `FastDLRATraceLoss`), where `λ_lat` has no effect. To get the latent term for a modality, pass ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` to the form with an external `x̂`. The latent term is zero for a weight-tied core with a right-invertible frontend, and need not be zero otherwise, including for a `CanonicalSortLayer` frontend and for a `CanonicalSlotLayer` frontend with more than one slot. Both forms of `CharbonnierRelativeLoss`, the cross-reconstruction term of `CrossModalCycleLoss` and the `α_rec` term of `IsometricRolloutLoss` compare fields in the form that `decode` returns, so a pre-split pair `(q = F⁺, p = F⁻)` of a `BilinearStiefelLayer` modality is recombined into the full field first. `IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.1, H = nothing)` evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty. It has no default latent Hamiltonian: `α_H > 0` requires `H`, and the constructor raises an `ArgumentError` without one, because a fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, `H` can be `nothing`, and the loss never calls it. `CrossModalCycleLoss(mod_a, mod_b; β_align = 1.0, β_cross = 1.0)` evaluates paired latent alignment ``\|z_a - z_b\|_F^2`` plus the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``. It has no cycle term, because ``\|\mathrm{encode}(nn, b, \mathrm{decode}(nn, b, z_a)) - z_a\|_F^2`` is zero for a weight-tied core with a right-invertible frontend and not in general. A negative or `NaN` weight raises an `ArgumentError`, and `x_a` and `x_b` whose latent states have different sizes, such as one sample against a batch, raise a `DimensionMismatch`. `IsometricRolloutLoss` raises an `ArgumentError` naming both element types when the data, the output of `integrator` or `H` has another element type than the network's parameters, and `CrossModalCycleLoss` does so for `x_a` or `x_b`. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (converted to the state element type, shifting the decoder lift by ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` and subtracting ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` before `rs.encoder` in `integrate_reduced_system`) and `pullback_mode::Symbol = :jacobian | :vjp`, where `:vjp` evaluates ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` in one reverse-mode pullback `Zygote.pullback(decoder, ξ)` and both modes use the preallocated buffers of `cache` (`ξ_buf`, `v_intermediate`, `f_intermediate`, and `jac_result` for `:jacobian`, with a per-element-type `IdDict{DataType, Any}` of caches for `ForwardDiff.Dual` Newton states). The default `cache` of `build_v_reduced` and `build_f_reduced` has the element type of the decoder's parameters, and the default `cache` of `HRedSys` has the element type of `ics.q`; a state of another floating-point element type than the cache's raises an `ArgumentError`. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. +- **`MultiModalSymplecticAutoencoder`, `encode`, `decode`, `translate`, `CharbonnierRelativeLoss`, `IsometricRolloutLoss`, and `CrossModalLoss` add multi-modal symplectic compression and multi-modal training losses, and `HRedSys` gains reference-state shifting, reverse-mode VJP vector fields, and preallocated work buffers.** All seven new names are exported. `MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression` combines a non-empty `NamedTuple` of modality frontends (individual layers or `Chain`s of symplectic particle layers) whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` (with `modality` a `Symbol` or `Val(modality)`) applies `frontends[modality]` followed by `encoder(core)`, `decode(nn, modality, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])`, and `translate(nn, from, to, x)` computes `decode(nn, to, encode(nn, from, x))`, raising an `ArgumentError` on an empty or dimension-mismatched `frontends` or an unknown modality symbol. `CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)` evaluates, for an external reconstruction `x̂`, the exact-zero Charbonnier relative reconstruction loss ``(\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon) / \sqrt{\|x\|_F^2 + \varepsilon^2} + \lambda_{\mathrm{lat}} \|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2 / (\|x\|_F^2 + \varepsilon^2)``, and without `x̂` the same formula for ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` (supporting low-rank `(X, S, V)` inputs via `FastDLRATraceLoss`). Both forms of `CharbonnierRelativeLoss`, the cross-reconstruction term of `CrossModalLoss` and the `α_rec` term of `IsometricRolloutLoss` compare fields in the form that `decode` returns, so a pre-split pair `(q = F⁺, p = F⁻)` of a `BilinearStiefelLayer` modality is recombined into the full field first. The two halves of a pre-split pair need one element type: halves of different element types raise an `ArgumentError` naming both types, with the message `A pre-split pair (q = F⁺, p = F⁻) needs one element type, got q of element type Float64 and p of element type Float32.` `IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.0, H = nothing)` evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty. `IsometricRolloutLoss(integrator)` has no energy term. `α_H > 0` requires `H`, and the constructor raises an `ArgumentError` without one, because a fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, a given `H` is not stored and never called. `CrossModalLoss(mod_a, mod_b; β_align = 1.0, β_cross = 1.0)` evaluates paired latent alignment ``\|z_a - z_b\|_F^2`` plus the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``. A negative or `NaN` weight raises an `ArgumentError`, and `x_a` and `x_b` whose latent states have different sizes, such as one sample against a batch, raise a `DimensionMismatch`. `IsometricRolloutLoss` raises an `ArgumentError` naming both element types when the data or the output of `integrator` has another element type than the network's parameters, or when `H` does so with `α_H > 0`, and `CrossModalLoss` does so for `x_a` or `x_b`. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (converted to the state element type, shifting the decoder lift by ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` and subtracting ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` before `rs.encoder` in `integrate_reduced_system`) and `pullback_mode::Symbol = :jacobian | :vjp`, where `:vjp` evaluates ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` in one reverse-mode pullback `Zygote.pullback(decoder, ξ)` and both modes use the preallocated buffers of `cache` (`ξ_buf`, `v_intermediate`, `f_intermediate`, and `jac_result` for `:jacobian`, with a per-element-type `IdDict{DataType, Any}` of caches for `ForwardDiff.Dual` Newton states). The default `cache` of `build_v_reduced` and `build_f_reduced` has the element type of the decoder's parameters, and the default `cache` of `HRedSys` has the element type of `ics.q`; a state of another floating-point element type than the cache's raises an `ArgumentError`. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. ### Changed diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index c9dfd43bf..99b13d2dd 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -76,24 +76,24 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. asserts the fields with `hasfield`. The fields stay because code outside this package may need them; removing them changes the type's parameters, its layout and those tests. -### K41 · `CrossModalCycleLoss(nn, x_a, x_b)` with a mixed-component `NamedTuple` raises a `MethodError`, not the documented `ArgumentError`. +### K41 · `CrossModalLoss(nn, x_a, x_b)` with a mixed-component `NamedTuple` raises a `MethodError`, not the documented `ArgumentError`. -- **location:** `src/loss/multimodal_losses.jl:385` +- **location:** `src/loss/multimodal_losses.jl:381` - **kind:** defect - **found:** #370 - **evidence:** - The three `(nn, x_a, x_b)` methods (lines 385, 391 and 397) annotate `x_a` and `x_b` with + The three `(nn, x_a, x_b)` methods (lines 381, 387 and 393) annotate `x_a` and `x_b` with `ArrayOrNamedTuple`, and AbstractNeuralNetworks defines that as `Union{AbstractArray{T}, NamedTupleOfArrays{T}}`, with one element type `T` for every component. A `NamedTuple` whose `q` and `p` differ in element type matches no `T`, so dispatch fails before - `_check_eltype` (line 254) runs. In a fresh `julia --startup-file=no --project=.` process: + `_check_eltype` (line 250) runs. In a fresh `julia --startup-file=no --project=.` process: ```julia T = typeof((q = zeros(Float32, 2, 3), p = zeros(Float64, 2, 3))) T <: AbstractNeuralNetworks.ArrayOrNamedTuple # false ``` - The pair form `CrossModalCycleLoss(nn, (x_a, x_b))` takes `Tuple{Any, Any}`, reaches + The pair form `CrossModalLoss(nn, (x_a, x_b))` takes `Tuple{Any, Any}`, reaches `_check_eltype`, and is tested. The docstring says that another element type raises an `ArgumentError`. Closing this entry means widening the `x_a` and `x_b` annotations of the three methods, for example to `Union{AbstractArray, NamedTuple}`. That is a design call about what the diff --git a/docs/src/reduced_order_modeling/losses.md b/docs/src/reduced_order_modeling/losses.md index 0abdaa6d1..86b5ae549 100644 --- a/docs/src/reduced_order_modeling/losses.md +++ b/docs/src/reduced_order_modeling/losses.md @@ -48,7 +48,7 @@ AutoEncoderLoss ReducedLoss CharbonnierRelativeLoss IsometricRolloutLoss -CrossModalCycleLoss +CrossModalLoss GeometricMachineLearning.ParametricLoss GeometricMachineLearning.apply_parametric projection_error diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 26873bd03..b5a186457 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -276,7 +276,7 @@ export LowRankStiefelLayer, DecodeLowRankStiefelLayer export FastDLRATraceLoss export SymmetricSymplecticAutoencoder export MultiModalSymplecticAutoencoder, encode, decode, translate -export CharbonnierRelativeLoss, IsometricRolloutLoss, CrossModalCycleLoss +export CharbonnierRelativeLoss, IsometricRolloutLoss, CrossModalLoss include("map_to_cpu.jl") diff --git a/src/loss/multimodal_losses.jl b/src/loss/multimodal_losses.jl index f017ac6ef..425401089 100644 --- a/src/loss/multimodal_losses.jl +++ b/src/loss/multimodal_losses.jl @@ -54,7 +54,17 @@ end # The data that a decoded state of a frontend is compared with, in the form that `decode` # returns for it. `decode` through a `BilinearStiefelLayer` returns the full field, so a pre-split # pair `(q = F⁺, p = F⁻)` is recombined into that field; every other input is compared as it is. +# The two halves of a pre-split pair must have one element type. +function _check_split_eltypes(q, p) + eltype(q) === eltype(p) || throw(ArgumentError( + "A pre-split pair (q = F⁺, p = F⁻) needs one element type, got q of element type $(eltype(q)) and p of element type $(eltype(p)).")) + nothing +end + +ChainRulesCore.@non_differentiable _check_split_eltypes(::Any, ::Any) + function _reconstruction_target(::BilinearStiefelLayer, x::NamedTuple{(:q, :p)}) + _check_split_eltypes(x.q, x.p) _parity_recombine(x.q, x.p) end _reconstruction_target(::Any, x) = x @@ -65,9 +75,7 @@ end @doc raw""" CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2) -Smooth Charbonnier relative reconstruction loss, with a latent consistency term when the reconstruction ``\hat{x}`` is passed. - -Called as `loss(nn, m, x̂, x)` with an external ``\hat{x}``, it evaluates +Smooth Charbonnier relative reconstruction and latent consistency loss: ```math \mathcal{L}_{\mathrm{Charb}}(\hat{x}, x) = @@ -75,16 +83,9 @@ Called as `loss(nn, m, x̂, x)` with an external ``\hat{x}``, it evaluates + \lambda_{\mathrm{lat}}\,\frac{\|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2}{\|x\|_F^2 + \varepsilon^2}. ``` -Called as `loss(nn, m, x)` (or `loss(nn, x)` for the first modality), it evaluates the reconstruction term alone, - -```math -\mathcal{L}_{\mathrm{Charb}}(x) = -\frac{\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon}{\sqrt{\|x\|_F^2 + \varepsilon^2}}, -\qquad \hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x)), -``` +Called as `loss(nn, m, x̂, x)`, it takes an external reconstruction ``\hat{x}``. Called as `loss(nn, m, x)` (or `loss(nn, x)` for the first modality), it uses ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))``. -Both forms compare ``\hat{x}`` with ``x`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality, passed as ``x`` or as ``\hat{x}``, is recombined into the full field. -The form without ``\hat{x}`` skips the latent term, so `λ_lat` has no effect there. The term is zero for a weight-tied core with a right-invertible frontend, because then ``\mathrm{encode}(nn, m, \hat{x}) = \mathrm{encode}(nn, m, x)``. It need not be zero for any other core, nor for a [`CanonicalSortLayer`](@ref) or [`CanonicalSlotLayer`](@ref) frontend, whose decoder does not map into the chamber of sorted states in general. To get the latent term there, pass ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` to the form with an external ``\hat{x}``. +Both forms compare ``\hat{x}`` with ``x`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality, passed as ``x`` or as ``\hat{x}``, is recombined into the full field. The two halves of such a pair have one element type; halves of different element types raise an `ArgumentError`. Here ``\varepsilon > 0`` (default `1e-3`) regularizes the Frobenius norm at ``\hat{x} = x`` and ``x = 0``, and ``\lambda_{\mathrm{lat}} \ge 0`` (default `1e-2`) weights the latent re-encoding consistency penalty. @@ -116,15 +117,6 @@ function _charbonnier_eval( (term_rec + term_lat)::T end -# The reconstruction term alone, with `x̂` compared with `target` and normalised by `x`. -function _charbonnier_rec_eval( - loss::CharbonnierRelativeLoss, x̂, target, x, ::Type{T}) where {T} - eps_T = T(loss.ε) - eps_sq = eps_T * eps_T - num_rec = sqrt(_frob_sq_diff(x̂, target) + eps_sq) - eps_T - (num_rec / sqrt(_frob_sq_norm(x) + eps_sq))::T -end - Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, x̂, x) T = promote_type(_state_eltype(x̂), _state_eltype(x), _leaf_eltype(ps.core)) @@ -147,8 +139,10 @@ end Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, x) T = promote_type(_state_eltype(x), _leaf_eltype(ps.core)) - x̂ = decode(model, ps, modality, encode(model, ps, modality, x)) - _charbonnier_rec_eval(loss, x̂, _reconstruction_target(model, modality, x), x, T) + z = encode(model, ps, modality, x) + x̂ = decode(model, ps, modality, z) + z_hat = encode(model, ps, modality, x̂) + _charbonnier_eval(loss, x̂, _reconstruction_target(model, modality, x), z_hat, z)::T end Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( @@ -178,7 +172,7 @@ Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( end @doc raw""" - IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.1, H = nothing) + IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.0, H = nothing) Multi-step latent rollout, decoded trajectory reconstruction, and latent Hamiltonian conservation loss: @@ -195,10 +189,10 @@ Multi-step latent rollout, decoded trajectory reconstruction, and latent Hamilto where ``z_k = \mathrm{encode}(nn, m, x_k)``, ``\hat{z}_1 = z_1``, ``\hat{z}_{k+1} = \mathrm{integrator}(\hat{z}_k)`` for ``k = 1, \dots, K - 1``, and `H` is a callable latent Hamiltonian. `H(ẑ)` takes the latent state as `encode` and `integrator` give it, a vector of length ``2n`` or a matrix of size ``(2n, B)`` for a batched snapshot, and returns one scalar, one value for a whole batch. The decoded state is compared with ``x_k`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality is recombined into the full field. -The weights default to `α_lat = 1.0`, `α_rec = 1.0`, and `α_H = 0.1`. +The weights default to `α_lat = 1.0`, `α_rec = 1.0`, and `α_H = 0.0`, so `IsometricRolloutLoss(integrator)` has no energy term. `H` has no default: the loss requires a latent Hamiltonian `H` when `α_H > 0`, and the constructor raises an `ArgumentError` without one. A fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, `H` can be `nothing`, and the loss then has no energy term: it stores `nothing` in place of a given `H` and never calls it. -The data, the output of `integrator` and the value of `H` have the element type of the network's parameters; another element type raises an `ArgumentError`. +The data, the output of `integrator` and, with `α_H > 0`, the value of `H` have the element type of the network's parameters; another element type raises an `ArgumentError`. """ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss integrator::IT @@ -211,7 +205,7 @@ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss integrator; α_lat::Real = 1.0, α_rec::Real = 1.0, - α_H::Real = 0.1, + α_H::Real = 0.0, H = nothing) (α_lat >= 0 && α_rec >= 0 && α_H >= 0) || throw(ArgumentError( "IsometricRolloutLoss requires non-negative weights α_lat, α_rec, α_H, got α_lat = $α_lat, α_rec = $α_rec, α_H = $α_H.")) @@ -323,7 +317,7 @@ Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( end @doc raw""" - CrossModalCycleLoss(mod_a::Symbol, mod_b::Symbol; β_align = 1.0, β_cross = 1.0) + CrossModalLoss(mod_a::Symbol, mod_b::Symbol; β_align = 1.0, β_cross = 1.0) Cross-modal paired latent alignment and cross-reconstruction loss between modalities `mod_a` and `mod_b`: @@ -335,24 +329,24 @@ Cross-modal paired latent alignment and cross-reconstruction loss between modali where ``z_a = \mathrm{encode}(nn, a, x_a)`` and ``z_b = \mathrm{encode}(nn, b, x_b)``. The pair ``(x_a, x_b)`` is the same physical state in two modalities: the alignment term asks that both encode to one latent state, and the cross-reconstruction term asks that the latent state of ``x_a`` decodes to ``x_b``. -The loss is not symmetric in ``a`` and ``b``: the cross-reconstruction term decodes only into modality `mod_b`, so `CrossModalCycleLoss(mod_b, mod_a)` is another loss. +The loss is not symmetric in ``a`` and ``b``: the cross-reconstruction term decodes only into modality `mod_b`, so `CrossModalLoss(mod_b, mod_a)` is another loss. `decode(nn, b, z_a)` is compared with ``x_b`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality is recombined into the full field. A batch is the sum of its samples' losses; `x_a` and `x_b` with different numbers of samples raise a `DimensionMismatch`. The weights are non-negative and default to `β_align = 1.0` and `β_cross = 1.0`; a negative or `NaN` weight raises an `ArgumentError`. `x_a` and `x_b` have the element type of the network's parameters; another element type raises an `ArgumentError`. """ -struct CrossModalCycleLoss{T <: Real} <: NetworkLoss +struct CrossModalLoss{T <: Real} <: NetworkLoss mod_a::Symbol mod_b::Symbol β_align::T β_cross::T - function CrossModalCycleLoss( + function CrossModalLoss( mod_a::Symbol, mod_b::Symbol; β_align::Real = 1.0, β_cross::Real = 1.0) (β_align >= 0 && β_cross >= 0) || throw(ArgumentError( - "CrossModalCycleLoss requires non-negative weights β_align and β_cross, got β_align = $β_align, β_cross = $β_cross.")) + "CrossModalLoss requires non-negative weights β_align and β_cross, got β_align = $β_align, β_cross = $β_cross.")) T = promote_type(typeof(float(β_align)), typeof(float(β_cross))) new{T}(mod_a, mod_b, T(β_align), T(β_cross)) end @@ -360,18 +354,18 @@ end function _check_sample_count(z_a, z_b) size(z_a) == size(z_b) || throw(DimensionMismatch( - "CrossModalCycleLoss: the latent states of x_a and x_b have different sizes, $(size(z_a)) and $(size(z_b)).")) + "CrossModalLoss: the latent states of x_a and x_b have different sizes, $(size(z_a)) and $(size(z_b)).")) nothing end ChainRulesCore.@non_differentiable _check_sample_count(::Any, ::Any) Base.@constprop :aggressive function _cross_modal_eval( - loss::CrossModalCycleLoss, model::MultiModalSymplecticLayer, + loss::CrossModalLoss, model::MultiModalSymplecticLayer, ps::NamedTuple, x_a, x_b) T = _leaf_eltype(ps.core) - _check_eltype("CrossModalCycleLoss", "x_a", x_a, T) - _check_eltype("CrossModalCycleLoss", "x_b", x_b, T) + _check_eltype("CrossModalLoss", "x_a", x_a, T) + _check_eltype("CrossModalLoss", "x_b", x_b, T) z_a = encode(model, ps, loss.mod_a, x_a) z_b = encode(model, ps, loss.mod_b, x_b) _check_sample_count(z_a, z_b) @@ -382,35 +376,35 @@ Base.@constprop :aggressive function _cross_modal_eval( res::T end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( model::MultiModalSymplecticLayer, ps::NamedTuple, x_a::ArrayOrNamedTuple, x_b::ArrayOrNamedTuple) _cross_modal_eval(loss, model, ps, x_a, x_b) end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( model::MultiModalSymplecticLayer, ps::NetworkParameters, x_a::ArrayOrNamedTuple, x_b::ArrayOrNamedTuple) _cross_modal_eval(loss, model, params(ps), x_a, x_b) end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, x_a::ArrayOrNamedTuple, x_b::ArrayOrNamedTuple) loss(nn.model, params(nn), x_a, x_b) end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( model::MultiModalSymplecticLayer, ps::NamedTuple, pair::Tuple{Any, Any}) _cross_modal_eval(loss, model, ps, pair[1], pair[2]) end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( model::MultiModalSymplecticLayer, ps::NetworkParameters, pair::Tuple{Any, Any}) _cross_modal_eval(loss, model, params(ps), pair[1], pair[2]) end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, pair::Tuple{Any, Any}) loss(nn.model, params(nn), pair) end diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 910923f1f..8526df5d5 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -278,10 +278,10 @@ end @testset "Multi-modal losses, zero-input edge case, exact-zero loss, and Zygote vs FD gradients" begin @test isdefined(GeometricMachineLearning, :CharbonnierRelativeLoss) @test isdefined(GeometricMachineLearning, :IsometricRolloutLoss) - @test isdefined(GeometricMachineLearning, :CrossModalCycleLoss) + @test isdefined(GeometricMachineLearning, :CrossModalLoss) @test CharbonnierRelativeLoss <: NetworkLoss @test IsometricRolloutLoss <: NetworkLoss - @test CrossModalCycleLoss <: NetworkLoss + @test CrossModalLoss <: NetworkLoss # Default keyword checks l_charb_def = CharbonnierRelativeLoss() @@ -290,27 +290,28 @@ end step_id = z -> z H_lat_def = z -> sum(abs2, z) / 2 - l_iso_def = IsometricRolloutLoss(step_id; H = H_lat_def) + l_iso_def = IsometricRolloutLoss(step_id) @test l_iso_def.α_lat == 1.0 @test l_iso_def.α_rec == 1.0 - @test l_iso_def.α_H == 0.1 - @test l_iso_def.H === H_lat_def + @test l_iso_def.α_H == 0.0 + @test l_iso_def.H === nothing + @test IsometricRolloutLoss(step_id; α_H = 0.1, H = H_lat_def).H === H_lat_def # There is no default latent Hamiltonian: α_H > 0 requires H, and α_H = 0 needs none. re_iso_H = r"^(?:ArgumentError: )?IsometricRolloutLoss requires a latent Hamiltonian H when α_H > 0, got α_H = 0\.1 and H = nothing\. Pass H, or set α_H = 0\.$" - @test_throws re_iso_H IsometricRolloutLoss(step_id) + @test_throws re_iso_H IsometricRolloutLoss(step_id; α_H = 0.1) @test_throws re_iso_H IsometricRolloutLoss(step_id; α_H = 0.1, H = nothing) @test IsometricRolloutLoss(step_id; α_H = 0.0).H === nothing # With α_H = 0 a given H is not stored, so it is never called. @test IsometricRolloutLoss(step_id; α_H = 0.0, H = H_lat_def).H === nothing - l_cyc_def = CrossModalCycleLoss(:particles, :grid) + l_cyc_def = CrossModalLoss(:particles, :grid) @test l_cyc_def.mod_a === :particles @test l_cyc_def.mod_b === :grid @test l_cyc_def.β_align == 1.0 @test l_cyc_def.β_cross == 1.0 @test !hasfield(typeof(l_cyc_def), :β_cycle) - @test_throws MethodError CrossModalCycleLoss(:particles, :grid; β_cycle = 0.1) + @test_throws MethodError CrossModalLoss(:particles, :grid; β_cycle = 0.1) # Edge validation for loss hyperparameters re_charb_eps = r"^(?:ArgumentError: )?CharbonnierRelativeLoss requires ε > 0 and λ_lat >= 0, got ε = 0\.0, λ_lat = 0\.01\.$" @@ -318,15 +319,15 @@ end @test_throws re_charb_eps CharbonnierRelativeLoss(; ε = 0.0) @test_throws re_charb_lam CharbonnierRelativeLoss(; λ_lat = -0.1) - re_iso = r"^(?:ArgumentError: )?IsometricRolloutLoss requires non-negative weights α_lat, α_rec, α_H, got α_lat = -1\.0, α_rec = 1\.0, α_H = 0\.1\.$" + re_iso = r"^(?:ArgumentError: )?IsometricRolloutLoss requires non-negative weights α_lat, α_rec, α_H, got α_lat = -1\.0, α_rec = 1\.0, α_H = 0\.0\.$" @test_throws re_iso IsometricRolloutLoss(step_id; α_lat = -1.0) - re_cross = r"^(?:ArgumentError: )?CrossModalCycleLoss requires non-negative weights β_align and β_cross, got β_align = 1\.0, β_cross = -0\.1\.$" - @test_throws re_cross CrossModalCycleLoss(:particles, :grid; β_cross = -0.1) - re_cross_nan = r"^(?:ArgumentError: )?CrossModalCycleLoss requires non-negative weights β_align and β_cross, got β_align = 1\.0, β_cross = NaN\.$" - @test_throws re_cross_nan CrossModalCycleLoss(:particles, :grid; β_cross = NaN) - re_align = r"^(?:ArgumentError: )?CrossModalCycleLoss requires non-negative weights β_align and β_cross, got β_align = -0\.1, β_cross = 1\.0\.$" - @test_throws re_align CrossModalCycleLoss(:particles, :grid; β_align = -0.1) + re_cross = r"^(?:ArgumentError: )?CrossModalLoss requires non-negative weights β_align and β_cross, got β_align = 1\.0, β_cross = -0\.1\.$" + @test_throws re_cross CrossModalLoss(:particles, :grid; β_cross = -0.1) + re_cross_nan = r"^(?:ArgumentError: )?CrossModalLoss requires non-negative weights β_align and β_cross, got β_align = 1\.0, β_cross = NaN\.$" + @test_throws re_cross_nan CrossModalLoss(:particles, :grid; β_cross = NaN) + re_align = r"^(?:ArgumentError: )?CrossModalLoss requires non-negative weights β_align and β_cross, got β_align = -0\.1, β_cross = 1\.0\.$" + @test_throws re_align CrossModalLoss(:particles, :grid; β_align = -0.1) @testset "Forward evaluation, @inferred, zero edge case, and exact-zero loss ($T)" for T in REAL_ELTYPES rng = Random.MersenneTwister(202) @@ -354,11 +355,11 @@ end expected_c = expected_c_rec + lam_T * sum(abs2, z_hat .- z) / (sum(abs2, F) + eps_T^2) - # The self-reconstruction form has no latent term; the form with an external x̂ has it. + # Both forms have the latent term; the self-reconstruction form uses x̂ = decode(encode(x)). val_c1 = @inferred loss_c(nn, :grid, F) val_c2 = @inferred loss_c(nn, :grid, F_hat, F) @test eltype(val_c1) === T && eltype(val_c2) === T - @test isapprox(val_c1, expected_c_rec; atol = 10 * eps(T), rtol = 10 * eps(T)) + @test isapprox(val_c1, expected_c; atol = 10 * eps(T), rtol = 10 * eps(T)) @test isapprox(val_c2, expected_c; atol = 10 * eps(T), rtol = 10 * eps(T)) # In the self-reconstruction form, a pre-split field (q = F⁺, p = F⁻) gives the value of @@ -372,6 +373,22 @@ end @test isapprox(@inferred(loss_c(nn, :grid, F_hat, F_split)), val_c2; atol = 10 * eps(T), rtol = 10 * eps(T)) @test @inferred(loss_c(nn, :grid, F_split, F_split)) === zero(T) + # The two halves of a pre-split pair have one element type: halves of different element + # types raise an ArgumentError in both forms, and a pair of one element type gives the + # value of the full field that it recombines to by hand. + if T === Float64 + F_split_mixed = (q = F_plus, p = Float32.(F_minus)) + re_mixed = r"^(?:ArgumentError: )?A pre-split pair \(q = F⁺, p = F⁻\) needs one element type, got q of element type Float64 and p of element type Float32\.$" + @test_throws re_mixed loss_c(nn, :grid, F_hat, F_split_mixed) + @test_throws re_mixed loss_c(nn, :grid, F_split_mixed) + F_minus_same = Float64.(F_split_mixed.p) + F_split_same = (q = F_plus, p = F_minus_same) + F_same_full = hcat((F_plus .+ F_minus_same) ./ sqrt(T(2)), + ((F_plus .- F_minus_same) ./ sqrt(T(2)))[:, end:-1:1]) + val_same = @inferred loss_c(nn, :grid, F_hat, F_split_same) + val_same_full = loss_c(nn, :grid, F_hat, F_same_full) + @test isapprox(val_same, val_same_full; rtol = 10 * eps(T)) + end # Independent perturbation F_pert != F_hat so z_hat != z and the latent penalty term is active (catches a dropped latent penalty) F_pert = F .+ T(0.5) .* ones(T, nx, nv) @@ -428,6 +445,13 @@ end @test isapprox(val_iso_noH, T(1.0) * err_lat + T(0.5) * err_rec; atol = 15 * eps(T), rtol = 15 * eps(T)) + # IsometricRolloutLoss(integrator) has the weights α_lat = α_rec = 1 and no energy term. + loss_iso_def = IsometricRolloutLoss(step_lin) + val_iso_def = @inferred loss_iso_def(nn, :grid, F_traj) + @test eltype(val_iso_def) === T + @test isapprox(val_iso_def, err_lat + err_rec; + atol = 15 * eps(T), rtol = 15 * eps(T)) + # A trajectory of pre-split fields gives the value of the trajectory of full fields. split_traj = [( q = (F_k[:, 1:(nv ÷ 2)] .+ F_k[:, nv:-1:(nv ÷ 2 + 1)]) ./ @@ -438,9 +462,9 @@ end @test isapprox(@inferred(loss_iso(nn, :grid, split_traj)), val_iso; atol = 15 * eps(T), rtol = 15 * eps(T)) - # CrossModalCycleLoss forward & @inferred: the alignment term and the + # CrossModalLoss forward & @inferred: the alignment term and the # cross-reconstruction term ‖decode(b, z_a) - x_b‖². - loss_cyc = CrossModalCycleLoss( + loss_cyc = CrossModalLoss( :particles, :grid; β_align = T(1.2), β_cross = T(0.3) ) x_p = (q = randn(rng, T, 1, 3), p = randn(rng, T, 1, 3)) @@ -453,9 +477,9 @@ end T(0.3) * sum(abs2, decode(nn, :grid, z_a) .- F) @test isapprox(val_cyc, expected_cyc; atol = 10 * eps(T), rtol = 10 * eps(T)) - # Reverse direction CrossModalCycleLoss(:grid, :particles): the particle decoder returns + # Reverse direction CrossModalLoss(:grid, :particles): the particle decoder returns # the flat state [vec(q); vec(p)]. - loss_cyc_rev = CrossModalCycleLoss( + loss_cyc_rev = CrossModalLoss( :grid, :particles; β_align = T(1.2), β_cross = T(0.5)) val_cyc_rev = @inferred loss_cyc_rev(nn, F, x_p) expected_cyc_rev = T(1.2) * sum(abs2, z_b .- z_a) + @@ -485,9 +509,9 @@ end @test (@inferred loss_cyc(nn, (x_p_b, F_b))) == val_cyc_b # x_a and x_b with different numbers of samples raise a DimensionMismatch. - re_samples = r"^(?:DimensionMismatch: )?CrossModalCycleLoss: the latent states of x_a and x_b have different sizes, \(2,\) and \(2, 3\)\.$" + re_samples = r"^(?:DimensionMismatch: )?CrossModalLoss: the latent states of x_a and x_b have different sizes, \(2,\) and \(2, 3\)\.$" @test_throws re_samples loss_cyc(nn, x_p, F_b) - re_samples_rev = r"^(?:DimensionMismatch: )?CrossModalCycleLoss: the latent states of x_a and x_b have different sizes, \(2, 3\) and \(2,\)\.$" + re_samples_rev = r"^(?:DimensionMismatch: )?CrossModalLoss: the latent states of x_a and x_b have different sizes, \(2, 3\) and \(2,\)\.$" @test_throws re_samples_rev loss_cyc(nn, x_p_b, F) # A pre-split field (q = F⁺, p = F⁻) for x_b gives the value of the full field F. The @@ -516,15 +540,16 @@ end # Another precision than the network's raises an ArgumentError that names both. S = T === Float32 ? Float64 : Float32 msg(name, what) = "$name: $what has element type $S, but the network has element type $T." - loss_iso_int = IsometricRolloutLoss(ξ -> S.(step_lin(ξ)); H = H_lat) + loss_iso_int = IsometricRolloutLoss(ξ -> S.(step_lin(ξ)); α_H = 0.1, H = H_lat) @test_throws msg("IsometricRolloutLoss", "the integrator's output") loss_iso_int( nn, :grid, F_traj) - loss_iso_H = IsometricRolloutLoss(step_lin; H = ξ -> S(0.5) * S(sum(abs2, ξ))) + loss_iso_H = IsometricRolloutLoss( + step_lin; α_H = 0.1, H = ξ -> S(0.5) * S(sum(abs2, ξ))) @test_throws msg("IsometricRolloutLoss", "the value of H") loss_iso_H( nn, :grid, F_traj) @test_throws msg("IsometricRolloutLoss", "the data") loss_iso( nn, :grid, [S.(F_k) for F_k in F_traj]) - @test_throws msg("CrossModalCycleLoss", "x_b") loss_cyc(nn, x_p, S.(F)) + @test_throws msg("CrossModalLoss", "x_b") loss_cyc(nn, x_p, S.(F)) # A particle state with one component of another precision than the network's raises # the same ArgumentError, whichever component it is. @@ -535,8 +560,8 @@ end nn, :particles, [x_mixed, x_p]) @test_throws msg("IsometricRolloutLoss", "the data") loss_iso( nn, :particles, [x_p, x_mixed]) - @test_throws msg("CrossModalCycleLoss", "x_a") loss_cyc(nn, (x_mixed, F)) - @test_throws msg("CrossModalCycleLoss", "x_b") loss_cyc_rev(nn, (F, x_mixed)) + @test_throws msg("CrossModalLoss", "x_a") loss_cyc(nn, (x_mixed, F)) + @test_throws msg("CrossModalLoss", "x_b") loss_cyc_rev(nn, (F, x_mixed)) end end @@ -569,7 +594,7 @@ end H_fn = ξ -> 0.5 * sum(abs2, ξ) + 0.05 * sum(ξ .^ 4) loss_iso = IsometricRolloutLoss( step_fn; α_lat = 1.0, α_rec = 0.8, α_H = 0.25, H = H_fn) - loss_cyc = CrossModalCycleLoss(:particles, :grid; β_align = 1.0, β_cross = 0.2) + loss_cyc = CrossModalLoss(:particles, :grid; β_align = 1.0, β_cross = 0.2) flat0, rebuild_ps = _flatten_params(nn.params) h = 1e-6 @@ -578,7 +603,7 @@ end for (name, f_ps) in ( ("CharbonnierRelativeLoss", ps -> loss_c(nn.model, ps, :grid, F)), ("IsometricRolloutLoss", ps -> loss_iso(nn.model, ps, :grid, F_traj)), - ("CrossModalCycleLoss", ps -> loss_cyc(nn.model, ps, x_p, F)) + ("CrossModalLoss", ps -> loss_cyc(nn.model, ps, x_p, F)) ) g_zyg = Zygote.gradient(f_ps, nn.params)[1] g_vec = _extract_grad_vec(g_zyg) @@ -821,14 +846,14 @@ end @test isapprox(vec(decode(nn, :b, z_a)), Dec * z_a; rtol = 10 * eps(T)) β_align, β_cross = T(0.7), T(1.3) - loss = CrossModalCycleLoss(:a, :b; β_align = β_align, β_cross = β_cross) + loss = CrossModalLoss(:a, :b; β_align = β_align, β_cross = β_cross) expected = β_align * sum(abs2, z_a - z_b) + β_cross * sum(abs2, Dec * z_a - vec(x_b)) val = @inferred loss(nn, x_a, x_b) @test val isa T @test isapprox(val, expected; rtol = 10 * eps(T)) # With β_align = 0 the loss is the cross term alone, which is not small and depends on x_b. - loss_cross = CrossModalCycleLoss(:a, :b; β_align = 0, β_cross = 1) + loss_cross = CrossModalLoss(:a, :b; β_align = 0, β_cross = 1) cross = sum(abs2, Dec * z_a - vec(x_b)) @test isapprox(loss_cross(nn, x_a, x_b), cross; rtol = 10 * eps(T)) @test cross > 1e-2 @@ -846,9 +871,9 @@ end cycle_sa = norm(encode(nn_sa, :b, decode(nn_sa, :b, z_a_sa)) - z_a_sa) @test cycle_sa > T(1e-3) * norm(z_a_sa) - # The self-reconstruction form of CharbonnierRelativeLoss has no latent term: on the untied - # core, where that term is not zero, the value is the reconstruction term alone. The form - # with an external x̂ = decode(encode(x)) keeps the latent term. + # The self-reconstruction form of CharbonnierRelativeLoss has the latent term: on the untied + # core, where that term is not zero, the value is the reconstruction term plus the latent + # term, the value of the form with an external x̂ = decode(encode(x)). ε = T(1e-2) loss_c = CharbonnierRelativeLoss(; ε = ε, λ_lat = 1) x̂_a = decode(nn_sa, :a, encode(nn_sa, :a, x_a)) @@ -857,7 +882,7 @@ end @test lat > T(1e-3) * rec val_self = @inferred loss_c(nn_sa, :a, x_a) @test val_self isa T - @test isapprox(val_self, rec; rtol = 10 * eps(T)) + @test isapprox(val_self, rec + lat; rtol = 10 * eps(T)) @test isapprox(loss_c(nn_sa, :a, x̂_a, x_a), rec + lat; rtol = 10 * eps(T)) end @@ -873,7 +898,7 @@ end documented = (CanonicalSortLayer, CanonicalSlotLayer, InverseCanonicalSlotLayer, SymmetricSymplecticAutoencoder, BilinearStiefelLayer, LowRankStiefelLayer, MultiModalSymplecticAutoencoder, CharbonnierRelativeLoss, IsometricRolloutLoss, - CrossModalCycleLoss) + CrossModalLoss) absent = ("collective canonical slots", "relative modes canonically", "canonical quadratic energy", "β_cycle", "symplectic or Poisson frontend") for X in documented @@ -891,7 +916,8 @@ end @test occursin("declare on the parity coordinates", docstring(X)) end @test occursin("requires a latent Hamiltonian", docstring(IsometricRolloutLoss)) - @test occursin("same physical state", docstring(CrossModalCycleLoss)) - @test occursin("not symmetric", docstring(CrossModalCycleLoss)) - @test occursin("skips", docstring(CharbonnierRelativeLoss)) + @test occursin("same physical state", docstring(CrossModalLoss)) + @test occursin("not symmetric", docstring(CrossModalLoss)) + @test !occursin("skips", docstring(CharbonnierRelativeLoss)) + @test occursin("have one element type", docstring(CharbonnierRelativeLoss)) end From e800bdc7e60354cb78bbcb45037e3723462613e2 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sat, 31 Oct 2026 09:17:29 +0100 Subject: [PATCH 5/6] =?UTF-8?q?Refuse=20a=20latent=20Hamiltonian=20H=20whe?= =?UTF-8?q?n=20=CE=B1=5FH=20=3D=200=20in=20IsometricRolloutLoss?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With the default α_H = 0, a passed H had no effect and was dropped without a word. The constructor now raises an ArgumentError that names α_H and H, and stores H as given otherwise. The docstring, the CHANGELOG bullet and the K41 line citations follow. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- KNOWN_ISSUES.md | 6 +++--- src/loss/multimodal_losses.jl | 9 +++++---- test/architectures/multimodal_symplectic_autoencoder.jl | 6 ++++-- 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b65227d5..0676ce3dc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -150,7 +150,7 @@ breaking release). - **`BilinearStiefelLayer`, `TransposeBilinearStiefelLayer`, `LowRankStiefelLayer`, `DecodeLowRankStiefelLayer`, `FastDLRATraceLoss`, and `solve!` on `BilinearStiefelLayer` add two-sided velocity-parity Stiefel frontends, a gauge-invariant dynamical low-rank Stiefel frontend and decoder, a fast low-rank trace loss, and Tucker-2 HOSVD initialization.** The five new types are exported. `BilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)` (`Nx >= 1`, even `Nv >= 2`, `1 <= mx <= Nx`, `1 <= mv <= Nv ÷ 2`) allocates `(Φx = StiefelManifold(...), Φv = StiefelManifold(...))` with ``\Phi_x \in \mathrm{St}(m_x, N_x)`` and ``\Phi_v \in \mathrm{St}(m_v, N_v / 2)``, splits an input array ``F`` of shape `(Nx, Nv)` or `(Nx, Nv, B)` into velocity-even ``F^+ = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} + F_{:, N_v:-1:N_v/2+1})`` and velocity-odd ``F^- = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} - F_{:, N_v:-1:N_v/2+1})`` components (or accepts a pre-split `(q, p)` `NamedTuple` of shape `(Nx, Nv ÷ 2)` or `(Nx, Nv ÷ 2, B)`), applies optional Cholesky mass factors `(Rx, Rv)` when `metric_factors` is supplied, and projects ``Q_{\mathrm{red}} = \Phi_x^T F^+ \Phi_v`` and ``P_{\mathrm{red}} = \Phi_x^T F^- \Phi_v`` to return vectorized `(q = vec(Q_red), p = vec(P_red))` of half-dimension ``n = m_x m_v``; `TransposeBilinearStiefelLayer` (`inverse_layer(l)`) is its exact right-inverse. `LowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)` takes a DLRA `NamedTuple` `(X = X, S = S, V = V)` of arbitrary rank ``r_{\mathrm{in}} \ge 1`` and evaluates ``Q_{\mathrm{red}} = (\Phi_x^T X) S ((V^+)^T \Phi_v)`` and ``P_{\mathrm{red}} = (\Phi_x^T X) S ((V^-)^T \Phi_v)`` in ``O(r_{\mathrm{in}}(m_x N_x + m_v N_v))`` operations without forming ``X S V^T``, matching `BilinearStiefelLayer` on ``F = X S V^T`` identically and remaining invariant under ``O(r_{\mathrm{in}}) \times O(r_{\mathrm{in}})`` gauge transformations, while `DecodeLowRankStiefelLayer` (`inverse_layer(l)`) computes the truncated SVD of ``[Q_{\mathrm{red}}\;\; P_{\mathrm{red}}]`` of rank ``r_{\mathrm{eff}} = \min(r_{\mathrm{out}}, m_x, 2 m_v)`` and QR re-orthonormalizes the lifted factors to return `(X = X_hat, S = S_hat, V = V_hat)` on the Stiefel manifold. `FastDLRATraceLoss()` evaluates ``\|X S V^T - \hat{X} \hat{S} \hat{V}^T\|_F^2`` in ``O(r^2(N_x + N_v))`` operations via Gram-matrix traces clamped at `zero(T)`. `solve!(nn, F_snapshots)` and `solve!(l, ps, F_snapshots)` initialize `Φx` and `Φv` on a `BilinearStiefelLayer` from the top-``m_x`` spatial and top-``m_v`` parity-velocity left singular vectors of `F_snapshots`. The `(q, p)` of the parity frontends is canonical only for the structure that they declare on the parity coordinates, not for the Vlasov bracket, and the docstrings say so. The tests are `test/layers/bilinear_and_low_rank_stiefel.jl`. -- **`MultiModalSymplecticAutoencoder`, `encode`, `decode`, `translate`, `CharbonnierRelativeLoss`, `IsometricRolloutLoss`, and `CrossModalLoss` add multi-modal symplectic compression and multi-modal training losses, and `HRedSys` gains reference-state shifting, reverse-mode VJP vector fields, and preallocated work buffers.** All seven new names are exported. `MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression` combines a non-empty `NamedTuple` of modality frontends (individual layers or `Chain`s of symplectic particle layers) whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` (with `modality` a `Symbol` or `Val(modality)`) applies `frontends[modality]` followed by `encoder(core)`, `decode(nn, modality, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])`, and `translate(nn, from, to, x)` computes `decode(nn, to, encode(nn, from, x))`, raising an `ArgumentError` on an empty or dimension-mismatched `frontends` or an unknown modality symbol. `CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)` evaluates, for an external reconstruction `x̂`, the exact-zero Charbonnier relative reconstruction loss ``(\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon) / \sqrt{\|x\|_F^2 + \varepsilon^2} + \lambda_{\mathrm{lat}} \|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2 / (\|x\|_F^2 + \varepsilon^2)``, and without `x̂` the same formula for ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` (supporting low-rank `(X, S, V)` inputs via `FastDLRATraceLoss`). Both forms of `CharbonnierRelativeLoss`, the cross-reconstruction term of `CrossModalLoss` and the `α_rec` term of `IsometricRolloutLoss` compare fields in the form that `decode` returns, so a pre-split pair `(q = F⁺, p = F⁻)` of a `BilinearStiefelLayer` modality is recombined into the full field first. The two halves of a pre-split pair need one element type: halves of different element types raise an `ArgumentError` naming both types, with the message `A pre-split pair (q = F⁺, p = F⁻) needs one element type, got q of element type Float64 and p of element type Float32.` `IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.0, H = nothing)` evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty. `IsometricRolloutLoss(integrator)` has no energy term. `α_H > 0` requires `H`, and the constructor raises an `ArgumentError` without one, because a fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, a given `H` is not stored and never called. `CrossModalLoss(mod_a, mod_b; β_align = 1.0, β_cross = 1.0)` evaluates paired latent alignment ``\|z_a - z_b\|_F^2`` plus the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``. A negative or `NaN` weight raises an `ArgumentError`, and `x_a` and `x_b` whose latent states have different sizes, such as one sample against a batch, raise a `DimensionMismatch`. `IsometricRolloutLoss` raises an `ArgumentError` naming both element types when the data or the output of `integrator` has another element type than the network's parameters, or when `H` does so with `α_H > 0`, and `CrossModalLoss` does so for `x_a` or `x_b`. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (converted to the state element type, shifting the decoder lift by ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` and subtracting ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` before `rs.encoder` in `integrate_reduced_system`) and `pullback_mode::Symbol = :jacobian | :vjp`, where `:vjp` evaluates ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` in one reverse-mode pullback `Zygote.pullback(decoder, ξ)` and both modes use the preallocated buffers of `cache` (`ξ_buf`, `v_intermediate`, `f_intermediate`, and `jac_result` for `:jacobian`, with a per-element-type `IdDict{DataType, Any}` of caches for `ForwardDiff.Dual` Newton states). The default `cache` of `build_v_reduced` and `build_f_reduced` has the element type of the decoder's parameters, and the default `cache` of `HRedSys` has the element type of `ics.q`; a state of another floating-point element type than the cache's raises an `ArgumentError`. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. +- **`MultiModalSymplecticAutoencoder`, `encode`, `decode`, `translate`, `CharbonnierRelativeLoss`, `IsometricRolloutLoss`, and `CrossModalLoss` add multi-modal symplectic compression and multi-modal training losses, and `HRedSys` gains reference-state shifting, reverse-mode VJP vector fields, and preallocated work buffers.** All seven new names are exported. `MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression` combines a non-empty `NamedTuple` of modality frontends (individual layers or `Chain`s of symplectic particle layers) whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` (with `modality` a `Symbol` or `Val(modality)`) applies `frontends[modality]` followed by `encoder(core)`, `decode(nn, modality, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])`, and `translate(nn, from, to, x)` computes `decode(nn, to, encode(nn, from, x))`, raising an `ArgumentError` on an empty or dimension-mismatched `frontends` or an unknown modality symbol. `CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)` evaluates, for an external reconstruction `x̂`, the exact-zero Charbonnier relative reconstruction loss ``(\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon) / \sqrt{\|x\|_F^2 + \varepsilon^2} + \lambda_{\mathrm{lat}} \|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2 / (\|x\|_F^2 + \varepsilon^2)``, and without `x̂` the same formula for ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` (supporting low-rank `(X, S, V)` inputs via `FastDLRATraceLoss`). Both forms of `CharbonnierRelativeLoss`, the cross-reconstruction term of `CrossModalLoss` and the `α_rec` term of `IsometricRolloutLoss` compare fields in the form that `decode` returns, so a pre-split pair `(q = F⁺, p = F⁻)` of a `BilinearStiefelLayer` modality is recombined into the full field first. The two halves of a pre-split pair need one element type: halves of different element types raise an `ArgumentError` naming both types, with the message `A pre-split pair (q = F⁺, p = F⁻) needs one element type, got q of element type Float64 and p of element type Float32.` `IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.0, H = nothing)` evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty. `IsometricRolloutLoss(integrator)` has no energy term. `α_H > 0` requires `H`, and the constructor raises an `ArgumentError` without one, because a fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0` the loss takes no `H`, and a given `H` raises an `ArgumentError` with the message `IsometricRolloutLoss: H has no effect when α_H = 0, got α_H = 0.0 and a latent Hamiltonian H. Set α_H > 0, or pass no H.` `CrossModalLoss(mod_a, mod_b; β_align = 1.0, β_cross = 1.0)` evaluates paired latent alignment ``\|z_a - z_b\|_F^2`` plus the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``. A negative or `NaN` weight raises an `ArgumentError`, and `x_a` and `x_b` whose latent states have different sizes, such as one sample against a batch, raise a `DimensionMismatch`. `IsometricRolloutLoss` raises an `ArgumentError` naming both element types when the data or the output of `integrator` has another element type than the network's parameters, or when `H` does so with `α_H > 0`, and `CrossModalLoss` does so for `x_a` or `x_b`. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (converted to the state element type, shifting the decoder lift by ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` and subtracting ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` before `rs.encoder` in `integrate_reduced_system`) and `pullback_mode::Symbol = :jacobian | :vjp`, where `:vjp` evaluates ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` in one reverse-mode pullback `Zygote.pullback(decoder, ξ)` and both modes use the preallocated buffers of `cache` (`ξ_buf`, `v_intermediate`, `f_intermediate`, and `jac_result` for `:jacobian`, with a per-element-type `IdDict{DataType, Any}` of caches for `ForwardDiff.Dual` Newton states). The default `cache` of `build_v_reduced` and `build_f_reduced` has the element type of the decoder's parameters, and the default `cache` of `HRedSys` has the element type of `ics.q`; a state of another floating-point element type than the cache's raises an `ArgumentError`. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. ### Changed diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 99b13d2dd..ebc725192 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -78,15 +78,15 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. ### K41 · `CrossModalLoss(nn, x_a, x_b)` with a mixed-component `NamedTuple` raises a `MethodError`, not the documented `ArgumentError`. -- **location:** `src/loss/multimodal_losses.jl:381` +- **location:** `src/loss/multimodal_losses.jl:382` - **kind:** defect - **found:** #370 - **evidence:** - The three `(nn, x_a, x_b)` methods (lines 381, 387 and 393) annotate `x_a` and `x_b` with + The three `(nn, x_a, x_b)` methods (lines 382, 388 and 394) annotate `x_a` and `x_b` with `ArrayOrNamedTuple`, and AbstractNeuralNetworks defines that as `Union{AbstractArray{T}, NamedTupleOfArrays{T}}`, with one element type `T` for every component. A `NamedTuple` whose `q` and `p` differ in element type matches no `T`, so dispatch fails before - `_check_eltype` (line 250) runs. In a fresh `julia --startup-file=no --project=.` process: + `_check_eltype` (line 251) runs. In a fresh `julia --startup-file=no --project=.` process: ```julia T = typeof((q = zeros(Float32, 2, 3), p = zeros(Float64, 2, 3))) diff --git a/src/loss/multimodal_losses.jl b/src/loss/multimodal_losses.jl index 425401089..6b90fcfd6 100644 --- a/src/loss/multimodal_losses.jl +++ b/src/loss/multimodal_losses.jl @@ -191,7 +191,7 @@ where ``z_k = \mathrm{encode}(nn, m, x_k)``, ``\hat{z}_1 = z_1``, ``\hat{z}_{k+1 The decoded state is compared with ``x_k`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality is recombined into the full field. The weights default to `α_lat = 1.0`, `α_rec = 1.0`, and `α_H = 0.0`, so `IsometricRolloutLoss(integrator)` has no energy term. -`H` has no default: the loss requires a latent Hamiltonian `H` when `α_H > 0`, and the constructor raises an `ArgumentError` without one. A fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, `H` can be `nothing`, and the loss then has no energy term: it stores `nothing` in place of a given `H` and never calls it. +`H` has no default: the loss requires a latent Hamiltonian `H` when `α_H > 0`, and the constructor raises an `ArgumentError` without one. A fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, the loss has no energy term and takes no `H`: a given `H` would have no effect, so the constructor raises an `ArgumentError`. The data, the output of `integrator` and, with `α_H > 0`, the value of `H` have the element type of the network's parameters; another element type raises an `ArgumentError`. """ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss @@ -211,10 +211,11 @@ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss "IsometricRolloutLoss requires non-negative weights α_lat, α_rec, α_H, got α_lat = $α_lat, α_rec = $α_rec, α_H = $α_H.")) (iszero(α_H) || H !== nothing) || throw(ArgumentError( "IsometricRolloutLoss requires a latent Hamiltonian H when α_H > 0, got α_H = $α_H and H = nothing. Pass H, or set α_H = 0.")) + (!iszero(α_H) || H === nothing) || throw(ArgumentError( + "IsometricRolloutLoss: H has no effect when α_H = 0, got α_H = $α_H and a latent Hamiltonian H. Set α_H > 0, or pass no H.")) T = promote_type(typeof(float(α_lat)), typeof(float(α_rec)), typeof(float(α_H))) - H_used = iszero(α_H) ? nothing : H - new{typeof(integrator), typeof(H_used), T}( - integrator, H_used, T(α_lat), T(α_rec), T(α_H)) + new{typeof(integrator), typeof(H), T}( + integrator, H, T(α_lat), T(α_rec), T(α_H)) end end diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 8526df5d5..4719cc45c 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -302,8 +302,10 @@ end @test_throws re_iso_H IsometricRolloutLoss(step_id; α_H = 0.1) @test_throws re_iso_H IsometricRolloutLoss(step_id; α_H = 0.1, H = nothing) @test IsometricRolloutLoss(step_id; α_H = 0.0).H === nothing - # With α_H = 0 a given H is not stored, so it is never called. - @test IsometricRolloutLoss(step_id; α_H = 0.0, H = H_lat_def).H === nothing + # With α_H = 0 a given H has no effect, so the constructor refuses it. + re_iso_H0 = r"^(?:ArgumentError: )?IsometricRolloutLoss: H has no effect when α_H = 0, got α_H = 0\.0 and a latent Hamiltonian H\. Set α_H > 0, or pass no H\.$" + @test_throws re_iso_H0 IsometricRolloutLoss(step_id; H = H_lat_def) + @test_throws re_iso_H0 IsometricRolloutLoss(step_id; α_H = 0.0, H = H_lat_def) l_cyc_def = CrossModalLoss(:particles, :grid) @test l_cyc_def.mod_a === :particles From cb2e3ea0be3e4e7381b2d9294a7b2427d15edbb2 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sat, 31 Oct 2026 09:22:22 +0100 Subject: [PATCH 6/6] Test that IsometricRolloutLoss treats alpha_H = -0.0 as zero A negative zero satisfies iszero, so the constructor refuses a given H with it and accepts it without one. Two tests pin both cases. Co-Authored-By: Claude Sonnet 5.5 --- test/architectures/multimodal_symplectic_autoencoder.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 4719cc45c..962e22b92 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -306,6 +306,10 @@ end re_iso_H0 = r"^(?:ArgumentError: )?IsometricRolloutLoss: H has no effect when α_H = 0, got α_H = 0\.0 and a latent Hamiltonian H\. Set α_H > 0, or pass no H\.$" @test_throws re_iso_H0 IsometricRolloutLoss(step_id; H = H_lat_def) @test_throws re_iso_H0 IsometricRolloutLoss(step_id; α_H = 0.0, H = H_lat_def) + # A negative zero is zero: it is refused with an H and accepted without one. + re_iso_H0_neg = r"^(?:ArgumentError: )?IsometricRolloutLoss: H has no effect when α_H = 0, got α_H = -0\.0 and a latent Hamiltonian H\. Set α_H > 0, or pass no H\.$" + @test_throws re_iso_H0_neg IsometricRolloutLoss(step_id; α_H = -0.0, H = H_lat_def) + @test IsometricRolloutLoss(step_id; α_H = -0.0).H === nothing l_cyc_def = CrossModalLoss(:particles, :grid) @test l_cyc_def.mod_a === :particles