Add StandardAutoencoder, an autoencoder of Dense layers

CHANGELOG.md

@@ -17,6 +17,7 @@ breaking release).
 
 ### Added
 
+- **`StandardAutoencoder` is an autoencoder without structure.** Its encoder and its decoder are feedforward networks of `Dense` layers: `n_encoder_layers` (resp. `n_decoder_layers`) layers with activation at the hidden width `width`, then an affine output layer. Nothing makes the decoder symplectic, and neither dimension has to be even. It is the unstructured counterpart of `SymplecticAutoencoder`, for comparing the two at a similar number of parameters. `encoder` and `decoder` split it as they split any `AutoEncoder`, and it saves and loads with `save` and `load(NeuralNetwork, path, arch)`. `test/architectures/standard_autoencoder_tests.jl` checks the split (also for `full_dim == reduced_dim`), the parameter count, that training reduces the loss and that the decoder is not symplectic at a random point; `test/integration/hdf5_support.jl` checks the save/load round trip.
 - **The test suite runs on an Apple GPU in its `metal` group.** `test/devices/metal.jl` checks the two device paths that were measured by hand and not tested before: the products of a `PoissonTensor` with a wrapped device array, `𝕁 * view(A, [1, 2, 3, 4], :)` and `𝕁 * B'`, which `ext/GPUArraysCoreExt.jl` carries; and the `tensor_mat_mul!` kernel on `MetalBackend()`. Each result must stay on the device and match the host in `Float32`, with scalar indexing off. `Pkg.test(test_args = ["metal"])` runs them without the subject tests, and a missing device records a skip, which happens inside a sandbox. Metal (1.10 or later) is a test dependency on every platform; it installs and precompiles on Linux and Windows, and nothing there loads it. On an Apple M4 Max with Metal 1.11.1 all five assertions pass.
 - **A `Metal` workflow runs the Metal tests on GitHub's `macos-15` runner**, for the `min` and `1` Julia versions. The runner's GPU is Apple's paravirtual device, which Metal.jl supports from 1.10 and on macOS 15 or later. The workflow includes a step "Metal is functional" before the tests that fails the job where the device is unavailable, so it cannot pass without having run the tests. The job is not a required check. However, the Metal tests also run in two required checks—CI.yml's `Julia min - macOS-latest - default` and `Julia 1 - macOS-latest - default`—which execute `test/runtests.jl` with no test arguments on Apple silicon, selecting the `metal` group.
 - **`test/quality/jet.jl` checks the kernel launchers with JET.** It runs in `core`, directly after `quality/aqua.jl`, and takes about 14 s. Each function that launches a kernel of `src/` on the `CPU()` backend, the pullbacks among them, and the `Batch` functor that `test/data_loader/batch_index_set.jl` measures with `@allocated` get one `JET.report_opt` line per element type that the tests use, with `target_modules = (GeometricMachineLearning,)`. Each `@kernel` method also gets a line that analyses its generated `cpu_<kernel>` function at the launcher's context, because JET drops the reports of a kernel body when it analyses the launcher. 142 lines pass on Julia 1.11.9 (JET 0.9.20) and 1.13.1 (JET 0.12.2). Three lines report a runtime dispatch on both versions and are `@test_broken` with issue #325: `PoissonTensor(::CPU, ::Int, ::DataType)`, and the input-and-output method of `convert_input_and_batch_indices_to_array` at `Float32` and `Float64`. Where JET does not work, the file records one `@test_skip`. JET joins `test/Project.toml` with no `[compat]` bound. `KNOWN_ISSUES.md` K2 records the Revise `EMFILE` messages that JET brings into the test log, and K3 the dispatches that the lines do not see: an unstable argument of a kernel launch, and two barriers in the Poisson-tensor kernel.

docs/src/reduced_order_modeling/pod_autoencoders.md

@@ -75,6 +75,8 @@ Main.proof(raw"If we take as test function ``\tilde{\psi}_i = (\nabla\Psi^\mathr
 " * Main.indentation * raw"and since this must be true for every ``i = 1, \ldots, n`` we obtain the desired expression for the reduced vector field.")
 ```
 
+A standard autoencoder of this kind, with `Dense` layers in the encoder and the decoder, is implemented as [`StandardAutoencoder`](@ref).
+
 Both POD and standard autoencoders suffer from the problem that they completely neglect the structure of the differential equation and the data they are applied to. This can have grave consequences [peng2016symplectic, tyranowski2023symplectic, buchfink2023symplectic](@cite). [Hamiltonian model order reduction](@ref "Hamiltonian Model Order Reduction") can improve the approximation significantly in these situations.
 
 ## Library Functions
@@ -85,6 +87,7 @@ GeometricMachineLearning.Encoder
 GeometricMachineLearning.Decoder
 GeometricMachineLearning.UnknownEncoder
 GeometricMachineLearning.UnknownDecoder
+StandardAutoencoder
 encoder
 decoder
 ```

src/GeometricMachineLearning.jl

@@ -233,6 +233,7 @@ include("architectures/sympnet.jl")
 include("architectures/autoencoder.jl")
 include("architectures/symplectic_autoencoder.jl")
 include("architectures/psd.jl")
+include("architectures/standard_autoencoder.jl")
 include("architectures/hamiltonian_neural_network.jl")
 include("architectures/lagrangian_neural_network.jl")
 include("architectures/transformer_neural_network.jl")
@@ -244,7 +245,7 @@ export LagrangianNeuralNetwork
 export SympNet, LASympNet, GSympNet
 export ClassificationTransformer, ClassificationLayer
 export VolumePreservingFeedForward
-export SymplecticAutoencoder, PSDArch
+export SymplecticAutoencoder, PSDArch, StandardAutoencoder
 export HamiltonianArchitecture, StandardHamiltonianArchitecture,
        GeneralizedHamiltonianArchitecture
 

src/architectures/standard_autoencoder.jl

@@ -0,0 +1,89 @@
+@doc raw"""
+    StandardAutoencoder(full_dim, reduced_dim)
+
+Make an instance of `StandardAutoencoder` for dimensions `full_dim` and `reduced_dim`.
+
+# The architecture
+
+A standard autoencoder [goodfellow2016deep](@cite) has an encoder ``\Psi^\mathrm{enc}:\mathbb{R}^N\to\mathbb{R}^n`` and a decoder ``\Psi^\mathrm{dec}:\mathbb{R}^n\to\mathbb{R}^N`` that are both feedforward networks of `Dense` layers, with no structure imposed on either. The encoder is
+
+```math
+\Psi^\mathrm{enc} = L^\mathrm{out}_{w\to{}n}\circ{}L_{w\to{}w}\circ\cdots\circ{}L_{w\to{}w}\circ{}L_{N\to{}w},
+```
+
+with `n_encoder_layers` nonlinear layers ``L``, the first of which maps ``\mathbb{R}^N`` to the hidden width ``w`` and the others ``\mathbb{R}^w`` to itself, followed by an affine output layer ``L^\mathrm{out}`` without activation. The decoder is built the same way from ``\mathbb{R}^n`` to ``\mathbb{R}^N`` with `n_decoder_layers` nonlinear layers.
+
+Unlike for a [`SymplecticAutoencoder`](@ref), ``\nabla\Psi^\mathrm{dec}`` is in general not symplectic, so the reduced system is not Hamiltonian, and neither ``N`` nor ``n`` has to be even. It is the unstructured counterpart of the [`SymplecticAutoencoder`](@ref), and serves to compare the two.
+
+# Arguments
+
+Besides the required arguments `full_dim` and `reduced_dim` you can provide the following keyword arguments:
+- `width::Integer = 2full_dim`: the hidden width ``w``.
+- `n_encoder_layers::Integer = 2`: the number of nonlinear layers in the encoder.
+- `n_decoder_layers::Integer = 2`: the number of nonlinear layers in the decoder.
+- `activation = tanh`: the activation of the nonlinear layers.
+
+# Examples
+
+The two parts compose to the whole network, and the number of parameters follows from the widths:
+
+```jldoctest
+using GeometricMachineLearning
+using GeometricMachineLearning: params
+
+arch = StandardAutoencoder(4, 2; width = 10, n_encoder_layers = 2, n_decoder_layers = 3)
+nn = NeuralNetwork(arch)
+x = rand(4)
+
+nn(x) ≈ decoder(nn)(encoder(nn)(x)), parameterlength(nn.model)
+
+# output
+
+(true, 476)
+```
+
+Here the encoder has ``(4 + 1)\cdot10 + (10 + 1)\cdot10 + (10 + 1)\cdot2 = 182`` parameters and the decoder ``(2 + 1)\cdot10 + 2\cdot(10 + 1)\cdot10 + (10 + 1)\cdot4 = 294``.
+"""
+struct StandardAutoencoder{AT} <: AutoEncoder
+    full_dim::Int
+    reduced_dim::Int
+    width::Int
+    n_encoder_layers::Int
+    n_decoder_layers::Int
+    n_encoder_blocks::Int
+    n_decoder_blocks::Int
+    activation::AT
+end
+
+function StandardAutoencoder(full_dim::Integer, reduced_dim::Integer;
+        width::Integer = 2full_dim,
+        n_encoder_layers::Integer = 2,
+        n_decoder_layers::Integer = 2,
+        activation = tanh)
+    @assert full_dim ≥ reduced_dim "The dimension of the full-order model has to be at least that of the reduced-order model!"
+    @assert n_encoder_layers ≥ 1 && n_decoder_layers ≥ 1 "Encoder and decoder need at least one nonlinear layer each!"
+    # The dimension changes once in each direction, from `full_dim` to `reduced_dim`; the hidden
+    # layers are not blocks in the sense of `compute_iterations`.
+    StandardAutoencoder{typeof(activation)}(full_dim, reduced_dim, width, n_encoder_layers,
+        n_decoder_layers, 2, 2, activation)
+end
+
+# The layers do not depend on intermediate dimensions, only on the two ends. `compute_iterations`
+# would build a range of step zero for `full_dim == reduced_dim`, which the constructor allows.
+compute_encoder_iterations(arch::StandardAutoencoder) = [arch.reduced_dim, arch.full_dim]
+compute_decoder_iterations(arch::StandardAutoencoder) = [arch.reduced_dim, arch.full_dim]
+
+function _standard_autoencoder_layers(arch::StandardAutoencoder, input_dim::Integer,
+        output_dim::Integer, n_layers::Integer)
+    (Dense(input_dim, arch.width, arch.activation),
+        Tuple(Dense(arch.width, arch.width, arch.activation) for _ in 2:n_layers)...,
+        Dense(arch.width, output_dim, identity))
+end
+
+function encoder_layers_from_iteration(arch::StandardAutoencoder, ::AbstractVector{<:Integer})
+    _standard_autoencoder_layers(arch, arch.full_dim, arch.reduced_dim, arch.n_encoder_layers)
+end
+
+function decoder_layers_from_iteration(arch::StandardAutoencoder, ::AbstractVector{<:Integer})
+    _standard_autoencoder_layers(arch, arch.reduced_dim, arch.full_dim, arch.n_decoder_layers)
+end

src/architectures/symplectic_autoencoder.jl

@@ -22,6 +22,8 @@ Because the decoder has this particular property, the reduced system can be desc
 
 where ``(\nabla_\xi\Psi^d)^+`` is the *symplectic inverse* of ``\nabla_\xi\Psi^d`` (for more details see the docs on the [`AutoEncoder`](@ref) type).
 
+For an autoencoder without this structure, whose decoder is not symplectic and whose reduced system is therefore not Hamiltonian, see [`StandardAutoencoder`](@ref).
+
 # Arguments
 
 Besides the required arguments `full_dim` and `reduced_dim` you can provide the following keyword arguments:

test/architectures/standard_autoencoder_tests.jl

@@ -0,0 +1,60 @@
+using GeometricMachineLearning
+using GeometricMachineLearning: params
+using Test
+using Zygote: jacobian
+import Random
+
+Random.seed!(123)
+
+function test_encoder_and_decoder(N::Integer, n::Integer; kwargs...)
+    nn = NeuralNetwork(StandardAutoencoder(N, n; kwargs...))
+    ae_encoder = encoder(nn)
+    ae_decoder = decoder(nn)
+
+    test_vec = rand(N)
+    test_mat = rand(N, N)
+
+    @test size(ae_encoder(test_vec)) == (n,)
+    @test size(ae_decoder(rand(n))) == (N,)
+    @test nn(test_vec) ≈ ae_decoder(ae_encoder(test_vec))
+    @test nn(test_mat) ≈ ae_decoder(ae_encoder(test_mat))
+end
+
+# Each nonlinear layer is one Dense layer with bias, and each part ends in an affine layer.
+function test_parameter_count(N::Integer, n::Integer, w::Integer, n_enc::Integer, n_dec::Integer)
+    nn = NeuralNetwork(StandardAutoencoder(N, n; width = w,
+        n_encoder_layers = n_enc, n_decoder_layers = n_dec))
+    n_encoder = (N + 1) * w + (n_enc - 1) * (w + 1) * w + (w + 1) * n
+    n_decoder = (n + 1) * w + (n_dec - 1) * (w + 1) * w + (w + 1) * N
+    @test parameterlength(nn.model) == n_encoder + n_decoder
+    @test parameterlength(encoder(nn).model) == n_encoder
+    @test parameterlength(decoder(nn).model) == n_decoder
+    @test length(encoder(nn).model.layers) == n_enc + 1
+    @test length(decoder(nn).model.layers) == n_dec + 1
+end
+
+function test_training_reduces_loss(N::Integer, n::Integer; n_epochs::Integer = 100)
+    dl = DataLoader(rand(N, 10 * N); autoencoder = true)
+    nn = NeuralNetwork(StandardAutoencoder(N, n))
+    loss = AutoEncoderLoss()
+    loss_before = loss(nn, dl.input)
+    o = Optimizer(Adam(), nn)
+    o(nn, dl, Batch(10), n_epochs; show_progress = false)
+    @test loss(nn, dl.input) < loss_before
+end
+
+# Nothing makes the decoder symplectic: ∇Ψᵀ 𝕁 ∇Ψ differs from 𝕁 at a random point.
+function test_not_symplectic(N::Integer, n::Integer)
+    ae_decoder = decoder(NeuralNetwork(StandardAutoencoder(N, n)))
+    J = jacobian(ae_decoder, rand(n))[1]
+    @test !(J' * PoissonTensor(N) * J ≈ PoissonTensor(n))
+end
+
+for (N, n) in ((10, 4), (4, 2), (7, 3), (4, 4))
+    test_encoder_and_decoder(N, n)
+    test_parameter_count(N, n, 2N, 2, 2)
+end
+test_encoder_and_decoder(4, 2; width = 16, n_encoder_layers = 1, n_decoder_layers = 4)
+test_parameter_count(4, 2, 40, 3, 4)
+test_training_reduces_loss(10, 4)
+test_not_symplectic(10, 4)

test/integration/hdf5_support.jl

@@ -78,6 +78,24 @@ end
     end
 end
 
+# Plain Dense layers; the architecture rebuilds the model that the weights are read into.
+@testset "save/load roundtrip: StandardAutoencoder" begin
+    arch = StandardAutoencoder(10, 4; width = 12, n_encoder_layers = 2, n_decoder_layers = 3)
+    nn = NeuralNetwork(arch)
+    x = rand(10)
+    y_before = nn(x)
+
+    mktempdir() do dir
+        path = joinpath(dir, "ae.h5")
+        save(path, nn)
+        nn2 = load(NeuralNetwork, path, arch)
+
+        @test _ps_eq(params(nn), params(nn2))
+        @test nn2(x) ≈ y_before
+        @test encoder(nn2)(x) ≈ encoder(nn)(x)
+    end
+end
+
 # LASympNet contains LinearLayer → SymmetricMatrix parameters.
 @testset "save/load roundtrip: LASympNet (SymmetricMatrix)" begin
     arch = LASympNet(4)

test/runtests.jl

@@ -46,6 +46,7 @@ if "core" in GROUPS
     @safetestset "SympNet integrator" include("architectures/sympnet_integrator.jl")
     @safetestset "PSD tests" include("architectures/psd_architecture_tests.jl")
     @safetestset "SymplecticAutoencoder tests" include("architectures/symplectic_autoencoder_tests.jl")
+    @safetestset "StandardAutoencoder tests" include("architectures/standard_autoencoder_tests.jl")
     @safetestset "Check if autoencoder error is lower than PSD error" include("architectures/sae_error_lower_than_psd_error.jl")
     @safetestset "Check reduced model" include("reduced_system/reduced_system.jl")
     @safetestset "_norm keeps the element type of its argument" include("norm_eltype.jl")