ImplicitLayers

ImplicitLayers

SciML's implicit layers — deep equilibrium networks and neural ODEs — as Lenticulum factors. The equilibrium family of Implicit Learners.md, alongside VariationalDiffusion.jl's diffusion family.

The thesis

DeepEquilibriumNetworks.jl and DiffEqFlux.jl both hand you an AbstractLuxLayer: a function $x \mapsto y$ whose direction is fixed when you construct it. But both are built out of a relation —

\[\text{DEQ:}\quad z = g_\theta(z,x) \qquad\qquad \text{NeuralODE:}\quad \frac{dz}{dt} = f_\theta(z,t)\]

— and the layer is that relation with one direction chosen and the solve sealed inside. So there are two ways to wrap them, and the difference is the whole point of this package:

LuxFactor works today on anything in the SciML ecosystem and asks nothing of it. The other two need the inner network rather than the assembled layer, and give back the bidirectionality the rest of the project is about.

The files

Concept notes: [[The Equilibrium Family]], [[DEQ as a Relation]], [[NeuralODE as an Invertible Factor]].

Two properties worth knowing up front

No SciML dependency, and no AD. Dependencies are LuxCore, Random, LinearAlgebra — the same choice VariationalDiffusion.jl makes and for the same reason. Solvers are derivative-free (Broyden is what DeepEquilibriumNetworks.jl uses anyway) and integrators are fixed-step explicit. Derivative information, where needed, comes from finite differences and is honestly labelled a test-scale tool: the real answer is a VJP from automatic differentiation. See [[solve]] §4.

Every inversion returns a DiracBelief. A root-find and an ODE solve both produce a point. Transporting a distribution would need GaussianBelief, which lives in the top-level Lenticulum package that no lib/ package may depend on — the third factor package in a row to hit this. See [[The Equilibrium Family]] §5.

abstract type AbstractIntegrator

A fixed-step explicit scheme. Deliberately not adaptive: an adaptive solver chooses different steps forwards and backwards, which destroys the round-trip property this package leans on (flow.md §3).

abstract type AbstractRootSolver

A method for solving $F(u) = 0$ given only the ability to evaluate $F$.

Derivative-freeness is a requirement, not a convenience: the residual of a DEQ contains a neural network, and this package has no automatic-differentiation dependency (see ImplicitLayers.md). It is also what DeepEquilibriumNetworks.jl does in practice — Broyden and limited-memory Broyden are its workhorses.

BroydenSolver(; maxiters = 100, tol = 1e-10, init_scale = 1.0)

"Good" Broyden: a quasi-Newton method maintaining an approximate inverse Jacobian updated by the Sherman–Morrison formula from secant pairs.

Derivative-free, and it is what DeepEquilibriumNetworks.jl reaches for. Note the initialisation: with $H_0 = I$ the first step is

\[u_1 = u_0 - H_0F(u_0) = u_0 - (u_0 - g(u_0)) = g(u_0)\]

— the first Broyden step is exactly a Picard step, and everything after it is the correction Picard never makes. Asserted in the test suite.

Stores a dense $n\times n$ matrix, so it is $O(n^2)$ in memory. Real DEQs use limited-memory Broyden for exactly this reason; see solve.md §4.2.

DEQFactor(cell, dims::NamedTuple; solver = BroydenSolver(), channels = (:x, :z),
          input = default_cell_input)

A deep equilibrium model as a relation: the fixed-point condition $z = g_\theta(z, x)$ presented as a residual, solvable for either channel.

cell is any LuxCore.AbstractLuxLayer computing $g_\theta(z,x)$ — including a DeepEquilibriumNetworks.jl cell, which is exactly such a layer. dims gives the two channel dimensions. Parameters and state are the cell's, untouched.

f = DEQFactor(my_cell, (x = 4, z = 4); solver = BroydenSolver(maxiters = 200))

Contrast LuxFactor, which wraps an already assembled DeepEquilibriumNetwork and gets one polarity — a Lux layer with extra steps.

DEQModel(factor, polarity)

The open model the polarity selects: "solve this fixed-point equation for the unobserved channel". Pure — a deterministic relation with no internal randomness.

EulerIntegrator(steps = 50)

Explicit Euler. First order, and present mainly as the thing RK4Integrator is checked against — its error is visible enough to make the order of accuracy testable.

LuxFactor(layer, :in => :out; dims = (nothing, nothing), postprocess = identity)

Any LuxCore.AbstractLuxLayer as a unidirectional factor.

The point of least resistance for reusing the SciML ecosystem: a DeepEquilibriumNetwork, a NeuralODE, a Chain, a Boltz vision backbone all satisfy the interface, so all of them wrap without this package knowing anything about them.

postprocess adapts the layer's output to a plain array — needed for the SciML layers whose call returns a solution object rather than a vector:

# DiffEqFlux's NeuralODE returns an ODESolution
LuxFactor(NeuralODE(net, (0.0, 1.0), Tsit5()), :x => :y;
          postprocess = sol -> Array(sol)[:, end])

# DeepEquilibriumNetworks returns the steady state directly, and reports the solve in `st`
LuxFactor(DeepEquilibriumNetwork(cell, NewtonRaphson()), :x => :z)

NeuralODEFactor(dynamics, dim; tspan = (0.0, 1.0), integrator = RK4Integrator(),
                channels = (:z0, :z1), input = default_dynamics_input)

A neural ODE as a bidirectional factor: the relation $\{(z_0,z_1) : z_1 = \Phi_{t_0\to t_1}(z_0)\}$ where $\Phi$ is the flow of $dz/dt = f_\theta(z,t)$.

dynamics is any LuxCore.AbstractLuxLayer computing $f_\theta(z,t)$ — the same object you would hand to DiffEqFlux.NeuralODE. Parameters and state are its own.

f = NeuralODEFactor(my_net, 4; tspan = (0.0, 1.0), integrator = RK4Integrator(steps = 100))

Both channels have dimension dim — a flow cannot change the dimension of its state, which is why this factor takes one dim rather than two and why a NeuralODE cannot be used as an encoder that compresses. Contrast DEQFactor, whose two channels are independent.

NeuralODEModel(factor, polarity)

The open model the polarity selects. Pure, and — unusually — exactly invertible, so the same object serves as its own inverse with the time direction flipped.

PicardSolver(; maxiters = 100, tol = 1e-10, damping = 1.0)

Damped fixed-point iteration $u \leftarrow u - \beta F(u)$.

For the DEQ residual $F(z) = z - g(z,x)$ this is exactly $z \leftarrow (1-\beta)z + \beta g(z,x)$: the naive "just keep applying the layer" loop.

It converges if and only if the iteration is a contraction, and nothing checks that in advance. That is the caveat Implicit Learners.md records — "this only works if the iteration converges, and unconstrained DEQs need not" — and the test suite exhibits a linear cell with spectral radius > 1 on which this solver diverges and BroydenSolver does not.

RK4Integrator(steps = 50)

Classical fourth-order Runge–Kutta, fixed step. The default: for the smooth non-stiff dynamics a NeuralODE usually has, steps = 50 of RK4 is far more accurate than anything Euler reaches, and it is Tsit5's spiritual ancestor without the adaptivity.

SolveReport(converged, iterations, residual, nfe)

What a root solve is willing to say about itself.

nfe (number of function evaluations) is reported because it is the honest cost measure for an implicit layer — DeepEquilibriumNetworks.jl reports it too, in its DeepEquilibriumSolution, for the same reason: iteration count means nothing when one Broyden step and one Picard step cost differently.

[!important] converged == false is not an error Bayesian Lens.md: a solver that stopped early is an inexact inversion, and inexact inversions are legal — the free energy records the cost. So the report is threaded out rather than thrown, and it is the caller's business what to do about it.

_divergence(vf, u, t) -> Real

$\operatorname{tr}(\partial \mathrm{vf}/\partial u)$ by forward differences.

apply_layer(f::LuxFactor, x, ps, st) -> (y, st)

One forward pass, with postprocess applied to the output.

cell_value(f, z, x, ps, st) -> (g, st)

One forward pass of the cell: $g_\theta(z, x)$.

default_cell_input(z, x) = (z, x)

How a DEQ cell is called: LuxCore.apply(cell, input(z, x), ps, st). Override via the input keyword for cells expecting a different arrangement — the factor deliberately knows nothing about the cell's internals.

default_dynamics_input(z, t) = (z, t)

How the dynamics layer is called: LuxCore.apply(dynamics, input(z, t), ps, st).

deq_sensitivity(f, x, z★, ps, st) -> ∂z★/∂x

The implicit function theorem at the solved fixed point: $\partial z^\ast/\partial x = (I - \partial_z g)^{-1}\partial_x g$.

Both Jacobians come from fd_jacobian, so this costs $O(\dim x + \dim z)$ forward passes and is a test-scale tool — the real answer is a VJP from automatic differentiation. It is here because it makes the IFT checkable: for a linear cell $g = Wz + Ux + b$ the exact answer is $(I-W)^{-1}U$, and the test suite compares against it.

This is also the quantity a Gaussian message would need in order to push a covariance through the layer, which deq.md §4.1 explains cannot currently be returned.

fd_jacobian(F, u; ε = 1e-7) -> Matrix

The dense Jacobian $\partial F/\partial u$ by forward differences: $n+1$ evaluations of F.

This is the package's only source of derivative information, and it is a test-scale tool. For a real DEQ the answer is a vector–Jacobian product from automatic differentiation, at $O(1)$ cost; here it is $O(n)$ forward passes and it loses roughly half the significant digits to the step size. It exists so that ift_sensitivity can be implemented and checked against a closed form, not so that it can be used at scale. See solve.md §4.1.

flow_forward(f, z₀, ps, st) -> (z₁, st)
flow_reverse(f, z₁, ps, st) -> (z₀, st)

Integrate $t_0 \to t_1$ and $t_1 \to t_0$ respectively. Same integrator, endpoints swapped — that is the entire implementation of the reverse direction.

flow_logdet(f, z₀, ps, st) -> (z₁, Δlogdet, st)

Forward flow together with $\int \operatorname{tr}\partial_z f_\theta\,dt$, so that

\[\log p(z_1) = \log p(z_0) - \Delta\!\log\!\det\]

the instantaneous change of variables. This is the piece that would make the factor transport a density rather than a point — and it cannot be used for that today, because the belief type it would need (GaussianBelief, or a normalising-flow belief) is not reachable from a lib/ package. See neuralode.md §4.2.

ift_sensitivity(Jz, Jx) -> Matrix

The implicit function theorem, applied to a fixed point $z^\ast = g(z^\ast, x)$:

\[(I - \partial_z g)\,\frac{\partial z^\ast}{\partial x} = \partial_x g \qquad\Longrightarrow\qquad \frac{\partial z^\ast}{\partial x} = (I - \partial_z g)^{-1}\,\partial_x g\]

One linear solve, no unrolling — the property that makes implicit layers $O(1)$ in memory, and the reason Implicit Learners.md lists the IFT as the equilibrium family's backward pass. See [[Backpropagation by the Implicit Function Theorem]].

Throws if $I - \partial_z g$ is singular. That is not a numerical accident: it is the statement that the fixed point is not locally unique, i.e. that the relation branches there. The algebraic family calls that locus the discriminant ([[Branches and the Discriminant]]); here it is the same phenomenon and the same failure.

integrate(vf, u₀, t₀, t₁, integrator) -> u₁

Integrate $du/dt = \mathrm{vf}(u, t)$ from t₀ to t₁.

t₁ < t₀ is allowed and is the whole point: dt is simply negative and the same scheme runs backwards. vf is called as vf(u, t).

integrate_with_divergence(vf, u₀, t₀, t₁, integrator) -> (u₁, Δlogdet)

Integrate the state and the log-density correction of a continuous normalising flow:

\[\frac{d}{dt}\log p(u(t)) = -\operatorname{tr}\frac{\partial \mathrm{vf}}{\partial u} \qquad\Longrightarrow\qquad \log p(u_1) = \log p(u_0) - \underbrace{\int_{t_0}^{t_1}\!\operatorname{tr}\,\partial_u \mathrm{vf}\,dt}_{\texttt{Δlogdet}}\]

This is the instantaneous change-of-variables of Chen et al. / FFJORD, and it is what turns a NeuralODE from a map on points into a map on densities — i.e. what would let this factor push a real belief rather than a Dirac.

[!warning] The trace is computed by a dense finite-difference Jacobian $O(n)$ evaluations of vf per step, so $O(n \cdot \mathrm{steps})$ in total. FFJORD's whole contribution is the Hutchinson estimator that avoids this; see flow.md §4.2. What is here is correct, checkable against a linear system, and unusable above small n.

residual(f::DEQFactor, x, z, ps, st) -> (r, st)

$r = z - g_\theta(z,x)$, the vector energy. Membership of the relation is $r \approx 0$, which is [[README]]'s definition of an implicit learner instantiated exactly.

residual(f::NeuralODEFactor, z₀, z₁, ps, st) -> (r, st)

$r = z_1 - \Phi_{t_0\to t_1}(z_0)$, the vector energy.

solve_input(f, z, ps, st; x₀) -> (x★, SolveReport, st)

Solve $z = g_\theta(z, x)$ for x, holding z fixed — the direction a DeepEquilibriumNetwork cannot be asked for.

Requires $\dim x = \dim z$. Even then it is a general nonlinear root-find in x with no contraction structure to lean on, so Broyden is the only sensible solver and convergence is a genuinely open question per call. That is why the SolveReport is returned rather than discarded.

solve_root(F, u₀, solver) -> (u, SolveReport)

Solve $F(u) = 0$ from the initial guess u₀.

F is called as F(u) and must return something the same shape as u. The residual norm that decides convergence is the Euclidean norm; tol is absolute, which is a simplification worth knowing about (solve.md §4.3).

solve_state(f, x, ps, st; z₀) -> (z★, SolveReport, st)

Solve $z = g_\theta(z, x)$ for z. The forward DEQ pass — what DeepEquilibriumNetworks.jl does.

z₀ defaults to zeros, which is SkipDeepEquilibriumNetwork's point of departure: a learned initial guess is strictly better and is not implemented (deq.md §4.4).

vectorfield(f, ps, st) -> (vf, stref)

The dynamics as a plain closure vf(z, t), plus the Ref accumulating the threaded Lux state.

The Ref is a wart with a cause: integrate wants a pure (u,t) -> du function, while LuxCore wants st threaded in and out of every call. Something has to hold the state across the integrator's inner calls, and a Ref is the smallest thing that does. See neuralode.md §4.1.

LenticulumCore.assemble(f::NeuralODEFactor, p, ps, st)

Produces a lens with ExactInversion — not SolverInversion.

That is a real claim, not an optimism: the inverse of a flow is the flow with negative time, and it is exact in exact arithmetic. What it is not is exact in floating point, and the gap is the integrator's round-trip error. neuralode.md §4.3 argues that ExactInversion is nonetheless the right label, and says what would change if you disagree.

LenticulumCore.energyspace(::DEQFactor)

$E_c = \mathbb{R}^{\dim z}$: the residual itself, as a vector.

This is what [[Scalar and Multivariate Energy]] asks for and what VariationalDiffusion's factor could not provide — here the residual is a genuine $\mathbb{R}^n$-valued object and the natural scalarisation is $\tfrac12\|\cdot\|^2$. Keeping it vector-valued is exactly what the implicit function theorem needs, per that note's §6.3.

LenticulumCore.invert(lens, π, inputs, ps, st) -> (DiracBelief, st)

Run the root solve in the direction the polarity chose.

Returns a DiracBelief: a root-find produces a point, not a distribution. Propagating uncertainty would need the linearisation of deq_sensitivity and a GaussianBelief to put it in — and GaussianBelief lives in the top-level Lenticulum package, which no lib/ package may depend on. See deq.md §4.1; the same wall is recorded in VariationalDiffusion's factor.md.

LenticulumCore.supported_polarities(f::DEQFactor)

Both directions when $\dim x = \dim z$; only the forward one otherwise.

The count is the point. A DeepEquilibriumNetwork has one direction, fixed when you construct it. This has two whenever the residual is square, and the second one is a genuinely new capability rather than a re-labelling — it solves the same equation for the other variable.

LenticulumCore.supported_polarities(f::NeuralODEFactor)

Both of them, always, with no dimension condition and no solver — because a flow is invertible by construction.

This is the only factor in the project for which both directions are equally cheap and equally exact. GaussianFactor is bidirectional but its two directions do different arithmetic (a pushforward vs a likelihood); DEQFactor's reverse direction is a root-find that may not converge; this one runs the identical integrator with the endpoints swapped.

Mycelium.factor_message(f::LuxFactor, target, polarity, inputs, prior, ps, st)

A DiracBelief on the output channel. Asking for a message on the input channel throws a PolarityError-shaped ArgumentError, because that is precisely the direction a Lux layer does not have.

Mycelium.factor_message(f::DEQFactor, target, polarity, inputs, prior, ps, st)

The factor → variable message: a DiracBelief on target, from a root solve.

The incoming prior is used as the solver's initial guess and nothing else. That is a genuinely good use for it — warm-starting from the current belief is what makes iterated message passing over implicit layers cheap — and it is not the same as conditioning on it. See deq.md §4.2 for why this message is a posterior rather than a likelihood.

Mycelium.factor_message(f::NeuralODEFactor, target, polarity, inputs, prior, ps, st)

A DiracBelief on target, obtained by integrating in the appropriate direction.

The prior is genuinely unused here — unlike DEQFactor, where it warm-starts the solver. A flow has nothing to warm-start.

Mycelium.local_free_energy(f::DEQFactor, msgs, ps, st)

$\tfrac12\|r\|^2$ at the incoming messages' points, or 0.0 when a channel carries no point.

Note what this is not: there is no entropy term, because the inversion returns a Dirac and a Dirac has none. So this factor contributes energy only, and the Bethe counting correction of [[Bethe Free Energy]] has nothing to correct here. deq.md §4.3.

Mycelium.local_free_energy(f::NeuralODEFactor, msgs, ps, st)

$\tfrac12\|z_1 - \Phi(z_0)\|^2$ when both endpoints carry points, else 0.0.

On the relation this is identically zero, because the flow is exact and both endpoints agree by construction. That makes it a poor diagnostic and a good invariant: a nonzero value means either the integrator is inaccurate or the two messages disagree, and the test suite uses it that way.