Deep Equilibrium Models

(Bai et al., 2019) introduced Discrete Deep Equilibrium Models which drives a Discrete Dynamical System to its steady-state. (Pal et al., 2022) extends this framework to Continuous Dynamical Systems which converge to the steady-stable in a more stable fashion. For a detailed discussion refer to (Pal et al., 2022).

To construct a continuous DEQ, any ODE solver compatible with DifferentialEquations.jl API can be passed as the solver. To construct a discrete DEQ, any root finding algorithm compatible with NonlinearSolve.jl API can be passed as the solver.

Choosing a Solver

Root Finding Algorithms

Using Root Finding Algorithms give fast convergence when possible, but these methods also tend to be unstable. If you must use a root finding algorithm, we recommend using:

NewtonRaphson or TrustRegion for small models
LimitedMemoryBroyden for large Deep Learning applications (with well-conditioned Jacobians)
NewtonRaphson(; linsolve = KrylovJL_GMRES()) for cases when Broyden methods fail

Note that Krylov Methods rely on efficient VJPs which are not available for all Lux models. If you think this is causing a performance regression, please open an issue in Lux.jl.

ODE Solvers

Using ODE Solvers give slower convergence, but are more stable. We generally recommend these methods over root finding algorithms. If you use implicit ODE solvers, remember to use Krylov linear solvers, see OrdinaryDiffEq.jl documentation for these. For most cases, we recommend:

VCAB3() for high tolerance problems
Tsit5() for high tolerance problems where VCAB3() fails
In all other cases, follow the recommendation given in OrdinaryDiffEq.jl documentation

Sensitivity Analysis

This package does not override SciMLSensitivity.jl's automatic sensealg choice. solve is called with sensealg = nothing unless a sensealg is passed through the layer constructor kwargs.

For the out-of-place SteadyStateProblems constructed here, SciMLSensitivity selects SteadyStateAdjoint(autodiff = false, autojacvec = ZygoteVJP()). Pass sensealg = SteadyStateAdjoint(; linsolve = LUFactorization()) for small models if the automatic linear solver is a poor fit.
For MultiScaleNeuralODE (functor parameters), SciMLSensitivity selects GaussAdjoint(; autojacvec = ZygoteVJP()). A faster alternative is BacksolveAdjoint(; autojacvec = ZygoteVJP()), with the usual stability caveats. See the SciMLSensitivity.jl documentation.

Public API

DeepEquilibriumNetworks.DeepEquilibriumNetworks — Module

DeepEquilibriumNetworks
DEQs

The DeepEquilibriumNetworks module. DEQs is an alias for the module.

Example

julia> DEQs === DeepEquilibriumNetworks
true

source

DeepEquilibriumNetworks.DeepEquilibriumNetwork — Type

DeepEquilibriumNetwork(model, solver; init = missing, jacobian_regularization = nothing,
    problem_type::Type = SteadyStateProblem{false}, kwargs...)

Deep Equilibrium Network as proposed in (Bai et al., 2019) and (Pal et al., 2022).

Arguments

model: Lux layer defining the equilibrium map.
solver: Solver for the equilibrium problem. ODE solvers and nonlinear solvers are both supported.

Keywords

init: Initial condition layer for the equilibrium problem. If nothing, the initial condition is set to zero(x). If missing, it is set to WrappedFunction(zero). Otherwise, pass a Lux layer called as init(x, ps, st).
jacobian_regularization: Jacobian stabilization backend. Supported values are nothing, AutoForwardDiff, AutoFiniteDiff, and AutoZygote.
problem_type: Equilibrium problem type. Use ODEProblem to construct an ODE-based network; defaults to SteadyStateProblem.
kwargs: Additional keyword arguments passed to SciMLBase.solve. Omitting sensealg uses SciMLSensitivity.jl's automatic adjoint choice.

Returns

Returns a Lux layer. Calling the layer returns the model output and a state whose solution field contains a DeepEquilibriumSolution.

Example

julia> model = DeepEquilibriumNetwork(
           Parallel(+, Dense(2, 2; use_bias=false), Dense(2, 2; use_bias=false)),
           NewtonRaphson(); verbose=false);

julia> rng = Xoshiro(0);

julia> ps, st = Lux.setup(rng, model);

julia> size(first(model(ones(Float32, 2, 1), ps, st)))
(2, 1)

source

DeepEquilibriumNetworks.DeepEquilibriumSolution — Type

DeepEquilibriumSolution(z_star, u0, residual, jacobian_loss, nfe, original)

Stores the solution data produced by a DeepEquilibriumNetwork and its variants.

Fields

z_star: Steady state, or the final iterate reached when the solver stops.
u0: Initial condition used by the equilibrium solve.
residual: Difference between $z^*$ and $f(z^*, x)$.
jacobian_loss: Jacobian stabilization loss.
nfe: Number of function evaluations.
original: Original solver solution object.

source

DeepEquilibriumNetworks.MultiScaleDeepEquilibriumNetwork — Method

MultiScaleDeepEquilibriumNetwork(main_layers::Tuple, mapping_layers::Matrix,
    post_fuse_layer::Union{Nothing, Tuple}, solver,
    scales::NTuple{N, NTuple{L, Int64}}; kwargs...)

Multi Scale Deep Equilibrium Network as proposed in (Bai et al., 2020).

Arguments

main_layers: Tuple of Neural Networks. Each Neural Network is applied to the corresponding scale.
mapping_layers: Matrix of Neural Networks. Each Neural Network is applied to the corresponding scale and the corresponding layer.
post_fuse_layer: Neural Network applied to the fused output of the main layers.
solver: Solver for the rootfinding problem. ODE Solvers and Nonlinear Solvers are both supported.
scales: Scales of the Multi Scale DEQ. Each scale is a tuple of integers. The length of the tuple is the number of layers in the corresponding main layer.

For keyword arguments, see DeepEquilibriumNetwork.

Example

julia> main_layers = (
           Parallel(+, Dense(4 => 4, tanh; use_bias=false), Dense(4 => 4, tanh; use_bias=false)),
           Dense(3 => 3, tanh), Dense(2 => 2, tanh), Dense(1 => 1, tanh));

julia> mapping_layers = [NoOpLayer() Dense(4 => 3, tanh) Dense(4 => 2, tanh) Dense(4 => 1, tanh);
                         Dense(3 => 4, tanh) NoOpLayer() Dense(3 => 2, tanh) Dense(3 => 1, tanh);
                         Dense(2 => 4, tanh) Dense(2 => 3, tanh) NoOpLayer() Dense(2 => 1, tanh);
                         Dense(1 => 4, tanh) Dense(1 => 3, tanh) Dense(1 => 2, tanh) NoOpLayer()];

julia> model = MultiScaleDeepEquilibriumNetwork(
           main_layers, mapping_layers, nothing, NewtonRaphson(), ((4,), (3,), (2,), (1,)));

julia> rng = Xoshiro(0);

julia> ps, st = Lux.setup(rng, model);

julia> x = rand(rng, Float32, 4, 12);

julia> size.(first(model(x, ps, st)))
((4, 12), (3, 12), (2, 12), (1, 12))

source

DeepEquilibriumNetworks.MultiScaleNeuralODE — Method

MultiScaleNeuralODE(args...; kwargs...)

Construct a multi-scale neural ODE with the same arguments as MultiScaleDeepEquilibriumNetwork.

This forwards all positional and keyword arguments to MultiScaleDeepEquilibriumNetwork, while setting problem_type to ODEProblem{false}.

Returns

Returns a MultiScaleDeepEquilibriumNetwork configured with ODE dynamics.

source

DeepEquilibriumNetworks.MultiScaleSkipDeepEquilibriumNetwork — Method

MultiScaleSkipDeepEquilibriumNetwork(main_layers::Tuple, mapping_layers::Matrix,
    post_fuse_layer::Union{Nothing, Tuple}, [init = nothing,] solver,
    scales::NTuple{N, NTuple{L, Int64}}; kwargs...)

Skip Multi Scale Deep Equilibrium Network as proposed in (Pal et al., 2022). Alias which creates a MultiScaleDeepEquilibriumNetwork with init kwarg set to passed value.

If init is not passed, it creates a MultiScale Regularized Deep Equilibrium Network.

Arguments

main_layers: Tuple of Lux layers, one per scale.
mapping_layers: Matrix of Lux layers mapping between scales.
post_fuse_layer: Optional tuple of Lux layers applied after scale fusion.
init: Optional tuple of Lux layers used to construct the initial conditions.
solver: Solver for the equilibrium problem.
scales: Output shape for each scale.

Returns

Returns a MultiScaleDeepEquilibriumNetwork.

source

DeepEquilibriumNetworks.SkipDeepEquilibriumNetwork — Method

SkipDeepEquilibriumNetwork(model, [init=nothing,] solver; kwargs...)

Skip Deep Equilibrium Network as proposed in (Pal et al., 2022).

This is a convenience constructor for DeepEquilibriumNetwork that forwards init through the init keyword argument. If init is omitted, the initial condition is nothing.

Arguments

model: Lux layer defining the equilibrium map.
init: Optional Lux layer used to construct the initial condition.
solver: Solver for the equilibrium problem.

Returns

Returns a DeepEquilibriumNetwork.

Example

julia> model = SkipDeepEquilibriumNetwork(
           Parallel(+, Dense(2, 2; use_bias=false), Dense(2, 2; use_bias=false)),
           NewtonRaphson(); verbose=false);

julia> ps, st = Lux.setup(Xoshiro(0), model);

julia> size(first(model(ones(Float32, 2, 1), ps, st)))
(2, 1)

source