SILVA Implicit DAE and Adversarial Residual Lab¶
This lab derives an implicit Runge-Kutta layer for index-1 differential-algebraic equations, validates one- and two-stage roots, and then adds an optional adversarial residual objective. DAE-PINNs motivate the implicit stage construction [52]. The similarly named DEQGAN uses “DEQ” to mean “Differential Equation,” not “Deep Equilibrium” [53].
Numbered literature: [1], [52], [53]. Each number opens the complete citation and its primary external source.
from pathlib import Path
import importlib.util
import subprocess
import sys
REPO_URL = "https://github.com/jseluis/silva-networks.git"
def find_local_silva_root():
candidates = [Path.cwd(), Path("/content/silva-networks")]
root = Path.cwd()
while root != root.parent:
candidates.append(root)
root = root.parent
for candidate in candidates:
if (candidate / "src" / "silva_networks").exists():
return candidate
return None
root = find_local_silva_root()
if root is not None:
sys.path.insert(0, str(root / "src"))
elif importlib.util.find_spec("silva_networks") is None:
subprocess.check_call([sys.executable, "-m", "pip", "install", f"git+{REPO_URL}"])
root = Path.cwd()
else:
root = Path.cwd()
import torch
import matplotlib.pyplot as plt
plt.rcParams.update({"figure.dpi": 300, "savefig.dpi": 300})
from silva_networks import (
SILVAImplicitDAEStep,
SILVAResidualDiscriminator,
make_linear_dae_dataset,
silva_adversarial_residual_loss,
)
torch.manual_seed(250)
<torch._C.Generator at 0x112068ed0>
1. Semi-Explicit Index-1 DAE¶
Separate differential and algebraic states:
$$ \dot y=f(y,z), \qquad 0=g(y,z). $$
The teaching system is
$$ \dot y=-y+z, \qquad g(y,z)=z-\frac12y. $$
Eliminating $z$ gives $\dot y=-y/2$, so $y(t)=y_0e^{-t/2}$ and $z(t)=y(t)/2$.
data = make_linear_dae_dataset(steps=8, dimensions=1, step_size=0.1)
assert data.constraint_residual().abs().max() == 0
print("trajectory:", tuple(data.differential.shape))
print("exact final differential state:", float(data.differential[-1]))
trajectory: (9, 1) exact final differential state: 0.3351600170135498
2. Implicit Runge-Kutta Stage Equations¶
For $s$ stages and Butcher coefficients $(A,b,c)$,
$$ Y_j=y_n+h\sum_{i=1}^s a_{ji}f(Y_i,Z_i), \qquad g(Y_j,Z_j)=0, $$
$$ y_{n+1}=y_n+h\sum_{i=1}^s b_i f(Y_i,Z_i), \qquad g(y_{n+1},z_{n+1})=0. $$
All stage and endpoint algebraic variables form one root vector. A damped Newton solve makes this root system an implicit SILVA layer. The dynamics may be known, learned, or a composition of both.
3. Backward Euler Is the One-Stage Case¶
With $A=[1]$, $b=[1]$, and $c=[1]$,
$$ y_{n+1}=y_n+h(-y_{n+1}+z_{n+1}), \qquad z_{n+1}=\frac12y_{n+1}. $$
Therefore
$$ y_{n+1}=\frac{y_n}{1+h/2}. $$
backward_euler = SILVAImplicitDAEStep(max_iter=6, tol=1e-8)
step = backward_euler(
data.differential[:1],
data.algebraic[:1],
data.step_size,
data.dynamics,
data.constraint,
)
discrete_exact = data.differential[:1] / (1.0 + data.step_size / 2.0)
assert step.residual < 1e-7
assert torch.allclose(step.differential, discrete_exact, atol=1e-6)
print("root residual:", step.residual)
print("differential/algebraic:", float(step.differential), float(step.algebraic))
root residual: 1.862645149230957e-09 differential/algebraic: 0.4761904776096344 0.2380952388048172
4. Two-Stage Gauss-Legendre Layer¶
The fourth-order two-stage tableau is
$$ A=\begin{bmatrix} 1/4 & 1/4-\sqrt3/6\\ 1/4+\sqrt3/6 & 1/4 \end{bmatrix}, \quad b=\begin{bmatrix}1/2&1/2\end{bmatrix}. $$
Changing the tableau changes the discretization while preserving the public DAE root contract.
root_three = 3.0**0.5
gauss = SILVAImplicitDAEStep(
a=torch.tensor([[0.25, 0.25-root_three/6], [0.25+root_three/6, 0.25]]),
b=torch.tensor([0.5, 0.5]),
c=torch.tensor([0.5-root_three/6, 0.5+root_three/6]),
max_iter=6,
tol=1e-8,
)
gauss_step = gauss(
data.differential[:1], data.algebraic[:1], data.step_size,
data.dynamics, data.constraint,
)
assert gauss_step.stage_differential.shape == (1, 2, 1)
assert gauss_step.residual < 1e-6
print("two-stage result:", float(gauss_step.differential))
two-stage result: 0.4756147265434265
5. Roll Out the Implicit Layer¶
Each step solves a local root. A trajectory is a sequence of such implicit points; it is not one globally weight-tied deep-equilibrium state. This classification keeps DAE time stepping distinct from a Bai-style DEQ while still placing the implicit layer inside SILVA.
y = data.differential[:1]
z = data.algebraic[:1]
trajectory = [y.detach().squeeze()]
root_residuals = []
for _ in range(8):
result = gauss(y, z, data.step_size, data.dynamics, data.constraint)
y, z = result.differential, result.algebraic
trajectory.append(y.detach().squeeze())
root_residuals.append(result.residual)
trajectory = torch.stack(trajectory)
print("maximum rollout root residual:", max(root_residuals))
maximum rollout root residual: 1.890951750738168e-08
figure, axes = plt.subplots(1, 2, figsize=(7.0, 2.6))
axes[0].plot(data.times[:, 0], data.differential[:, 0], label="continuous exact")
axes[0].plot(data.times[:, 0], trajectory, "--", label="implicit rollout")
axes[0].set(xlabel="time", ylabel="differential state")
axes[0].legend()
axes[1].semilogy(root_residuals, marker="o", markersize=3)
axes[1].set(xlabel="time step", ylabel="root residual")
figure.tight_layout()
plt.show()
6. Adversarial Residual Objective and Naming Boundary¶
An optional discriminator can compare equation residuals $r_\theta$ with a near-zero reference distribution $r_0$. With discriminator $D_\omega$,
$$ \mathcal L_D =-\mathbb E\log D_\omega(r_0) -\mathbb E\log(1-D_\omega(r_\theta)), $$
$$ \mathcal L_G=-\mathbb E\log D_\omega(r_\theta). $$
This objective follows the differential-equation GAN work [53]. It is not an
equilibrium solver and is deliberately exposed as
silva_adversarial_residual_loss, not as a selectable DEQ family.
discriminator = SILVAResidualDiscriminator(residual_dim=1, hidden_dim=8, depth=1)
equation_residual = data.constraint(data.differential, data.algebraic)
equation_residual = equation_residual + 0.02 * torch.randn_like(equation_residual)
losses = silva_adversarial_residual_loss(
discriminator,
equation_residual,
reference=torch.zeros_like(equation_residual),
instance_noise=0.005,
)
assert torch.isfinite(losses.generator + losses.discriminator)
print("generator/discriminator:", float(losses.generator), float(losses.discriminator))
generator/discriminator: 0.6628165245056152 1.3866183757781982
7. Choosing the Physics Construction¶
Use a physics-informed equilibrium when the solution representation itself is a fixed point and implicit time derivatives are needed. Use an implicit DAE step when every time advance must satisfy differential and algebraic stage equations. Add adversarial residual matching only when distributional residual training is part of the study; it supplements rather than replaces direct constraint, root, and trajectory diagnostics.
From 25 Silva Implicit Dae And Residuals to a Custom SILVA Family¶
The construction in this notebook can be separated into the universal conditioned-equilibrium contract
$$ z_0=I_\eta(x),\qquad z^\star=T_\theta(z^\star,x),\qquad \widehat y=Q_\psi(z^\star). $$
For this topic:
| Part | Concrete interpretation |
|---|---|
| Equilibrium state | an implicit latent state or coupled differential/algebraic stage state |
| Condition | time, initial/boundary values, dynamics, and algebraic constraints |
| Repeated computation | a time-conditioned fixed point or implicit Runge-Kutta root map |
| Required invariants | initial/boundary conditions, equation shape, and constraint consistency |
| Replaceable components | time/source lift, transition, readout, dynamics, constraints, losses, and solvers |
The initializer and source path are evaluated outside or alongside the root solve. Only the state-preserving transition is repeated. Replacing an internal architecture does not change this equation, provided the transition still maps the same state space into itself.
Supply a New DAE and Implicit Tableau¶
step = SILVAImplicitDAEStep(
a=runge_kutta_a,
b=runge_kutta_b,
c=runge_kutta_c,
linear_solver="gmres",
linear_max_iter=100,
linear_tol=1e-7,
)
result = step(
differential_state,
algebraic_state,
step_size,
dynamics=my_differential_field,
constraint=my_algebraic_constraint,
)
The dimensions of the differential and algebraic states may change between applications, while each supplied dynamics and constraint function must preserve its declared equation shape.
import torch as silva_extension_torch
from torch import nn as silva_extension_nn
from silva_networks import (
SILVAConditionedEquilibrium,
SILVAZeroInitializer,
SolverConfig,
validate_silva_transition,
)
class NotebookExtensionTransition(silva_extension_nn.Module):
def __init__(self, condition_dim=2, state_dim=3):
super().__init__()
self.source = silva_extension_nn.Linear(condition_dim, state_dim)
self.state_field = silva_extension_nn.Sequential(
silva_extension_nn.Linear(state_dim, 2 * state_dim),
silva_extension_nn.Tanh(),
silva_extension_nn.Linear(2 * state_dim, state_dim),
)
def forward(self, state, condition):
return silva_extension_torch.tanh(
self.source(condition) + 0.15 * self.state_field(state)
)
silva_extension_torch.manual_seed(610)
notebook_condition = silva_extension_torch.linspace(-1.0, 1.0, 8).reshape(4, 2)
notebook_state0 = silva_extension_torch.zeros(4, 3)
notebook_transition = NotebookExtensionTransition()
notebook_report = validate_silva_transition(
notebook_transition,
notebook_state0,
notebook_condition,
)
assert notebook_report.valid
with silva_extension_torch.no_grad():
notebook_reference_step = silva_extension_torch.tanh(
notebook_transition.source(notebook_condition)
+ 0.15 * notebook_transition.state_field(notebook_state0)
)
silva_extension_torch.testing.assert_close(
notebook_transition(notebook_state0, notebook_condition),
notebook_reference_step,
)
notebook_custom_model = SILVAConditionedEquilibrium(
notebook_transition,
SILVAZeroInitializer(3),
readout=silva_extension_nn.Linear(3, 1),
config=SolverConfig(
solver="picard",
max_iter=40,
tol=1e-7,
backward_mode="implicit",
backward_solver="gmres",
anderson_batch_dims=1,
),
)
notebook_custom_result = notebook_custom_model(
notebook_condition,
return_result=True,
)
assert notebook_custom_result.output.shape == (4, 1)
assert notebook_custom_result.solver_result.residual < 1e-5
notebook_custom_result.output.square().mean().backward()
assert all(
parameter.grad is not None and silva_extension_torch.isfinite(parameter.grad).all()
for parameter in notebook_custom_model.parameters()
)
print("custom transition:", notebook_report)
print("equilibrium residual:", notebook_custom_result.solver_result.residual)
custom transition: SILVATransitionReport(state_shape=(4, 3), output_shape=(4, 3), preserves_shape=True, preserves_device=True, preserves_dtype=True, finite=True, differentiable=True, parameter_count=54) equilibrium residual: 5.960464477539063e-08
Numerical Equivalence, Compact Reproduction, and Scale¶
Before training, compare one packaged transition with an independently written update:
$$ e_{\mathrm{step}} =\frac{\|T_\theta(z,x)-T_{\mathrm{ref}}(z,x)\|_2} {\|T_{\mathrm{ref}}(z,x)\|_2+\varepsilon}. $$
After solving, report the fixed-point residual separately:
$$ e_{\mathrm{fp}} =\frac{\|T_\theta(z^\star,x)-z^\star\|_2} {\|z^\star\|_2+\varepsilon}. $$
For this notebook, a compact reproduction must declare and assert equation residual, boundary error, trajectory error, and solver residual. A full experiment must additionally record the source dataset version and split, preprocessing, architecture widths, solver and optimizer schedules, random seeds, baseline configuration, checkpoints, and every deviation from the cited protocol.
The principal scaling axes are collocation count, latent dimension, stages, stiffness, and time horizon. Increase one axis at a time, retain the compact deterministic case as a regression test, and record task error, domain-specific residual, forward residual, backward linear residual, memory use, and runtime independently.
Extension Exercises¶
- Replace one component from this notebook while preserving its state and domain invariants.
- Write the replacement first as an independent reference function, then as a module, and assert one-step equivalence.
- Compare two solver configurations on the identical trained transition.
- Add a compact baseline and a predeclared metric threshold.
- Create a full-scale configuration without weakening the compact tests.
The complete authoring protocol is documented in Extending SILVA.
notebook_reproduction_record = {
"notebook": '25_silva_implicit_dae_and_residuals.ipynb',
"state": 'an implicit latent state or coupled differential/algebraic stage state',
"condition": 'time, initial/boundary values, dynamics, and algebraic constraints',
"transition": 'a time-conditioned fixed point or implicit Runge-Kutta root map',
"invariants": 'initial/boundary conditions, equation shape, and constraint consistency',
"compact_metric": 'equation residual, boundary error, trajectory error, and solver residual',
"scale_axis": 'collocation count, latent dimension, stages, stiffness, and time horizon',
}
assert all(notebook_reproduction_record.values())
notebook_reproduction_record
{'notebook': '25_silva_implicit_dae_and_residuals.ipynb',
'state': 'an implicit latent state or coupled differential/algebraic stage state',
'condition': 'time, initial/boundary values, dynamics, and algebraic constraints',
'transition': 'a time-conditioned fixed point or implicit Runge-Kutta root map',
'invariants': 'initial/boundary conditions, equation shape, and constraint consistency',
'compact_metric': 'equation residual, boundary error, trajectory error, and solver residual',
'scale_axis': 'collocation count, latent dimension, stages, stiffness, and time horizon'}
Worked Convergence and Sensitivity Study¶
The preceding example demonstrates one configured solve. This additional study changes the implicit dynamics factor while keeping the source fixed, so solver effort and implicit sensitivity can be read separately from task behavior. Locally, one eigendirection of a nonlinear transition can be represented by
$$ z_{k+1} = \rho z_k + u, \qquad 0 \leq \rho < 1. $$
Its equilibrium is
$$ z^\star = \frac{u}{1-\rho}. $$
Subtracting the fixed-point equation from the iteration gives the exact error recursion
$$ e_{k+1} = \rho e_k, \qquad |e_k| = \rho^k |e_0|. $$
For a requested absolute tolerance $\tau$, the idealized iteration estimate is
$$ k \geq \frac{\log(\tau/|e_0|)}{\log \rho}. $$
The same factor controls sensitivity. Differentiating the equilibrium with respect to the source gives
$$ \frac{\partial z^\star}{\partial u} =\frac{1}{1-\rho}. $$
Thus a transition can remain contractive while becoming expensive and highly sensitive as $\rho$ approaches one. The table and figure below measure this effect rather than merely stating it. They provide a reference envelope for the notebook's actual state, an implicit latent state or coupled differential/algebraic stage state, and its repeated map, a time-conditioned fixed point or implicit Runge-Kutta root map. The scalar study does not replace the domain model; it supplies a result whose convergence rate and derivative are known exactly, so the same reporting code can be trusted before it is applied to the larger transition.
import math as silva_deepening_math
import torch as silva_deepening_torch
silva_deepening_rates = (0.20, 0.45, 0.70, 0.85)
silva_deepening_source = 0.35
silva_deepening_tolerance = 1e-8
silva_deepening_histories = {}
silva_deepening_rows = []
for silva_deepening_rho in silva_deepening_rates:
silva_deepening_state = silva_deepening_torch.tensor(0.0)
silva_deepening_exact = silva_deepening_source / (1.0 - silva_deepening_rho)
silva_deepening_history = []
for silva_deepening_iteration in range(1, 241):
silva_deepening_next = (
silva_deepening_rho * silva_deepening_state + silva_deepening_source
)
silva_deepening_residual = abs(
float(silva_deepening_next - silva_deepening_state)
)
silva_deepening_history.append(silva_deepening_residual)
silva_deepening_state = silva_deepening_next
if silva_deepening_residual < silva_deepening_tolerance:
break
silva_deepening_u = silva_deepening_torch.tensor(
silva_deepening_source, requires_grad=True
)
silva_deepening_solution = silva_deepening_u / (1.0 - silva_deepening_rho)
silva_deepening_solution.backward()
silva_deepening_expected_sensitivity = 1.0 / (1.0 - silva_deepening_rho)
silva_deepening_gradient_error = abs(
float(silva_deepening_u.grad) - silva_deepening_expected_sensitivity
)
silva_deepening_histories[silva_deepening_rho] = silva_deepening_history
silva_deepening_rows.append(
(
silva_deepening_rho,
silva_deepening_iteration,
silva_deepening_history[-1],
abs(float(silva_deepening_state) - silva_deepening_exact),
float(silva_deepening_u.grad),
silva_deepening_gradient_error,
)
)
print('implicit dynamics factor')
print("rho | iterations | final residual | exact-state error | sensitivity | gradient error")
for silva_deepening_row in silva_deepening_rows:
print(
f"{silva_deepening_row[0]:.2f} | {silva_deepening_row[1]:3d} | "
f"{silva_deepening_row[2]:.3e} | {silva_deepening_row[3]:.3e} | "
f"{silva_deepening_row[4]:.4f} | {silva_deepening_row[5]:.3e}"
)
assert all(row[2] < silva_deepening_tolerance for row in silva_deepening_rows)
assert all(row[3] < 1e-6 for row in silva_deepening_rows)
assert all(row[5] < 1e-6 for row in silva_deepening_rows)
implicit dynamics factor rho | iterations | final residual | exact-state error | sensitivity | gradient error 0.20 | 12 | 0.000e+00 | 5.551e-17 | 1.2500 | 0.000e+00 0.45 | 23 | 0.000e+00 | 1.084e-08 | 1.8182 | 6.502e-08 0.70 | 45 | 0.000e+00 | 1.589e-07 | 3.3333 | 7.947e-08 0.85 | 93 | 0.000e+00 | 5.563e-07 | 6.6667 | 1.589e-07
import matplotlib.pyplot as silva_deepening_plt
silva_deepening_plt.rcParams.update({"figure.dpi": 300, "savefig.dpi": 300})
silva_deepening_figure, silva_deepening_axes = silva_deepening_plt.subplots(
1, 2, figsize=(8.6, 3.2)
)
for silva_deepening_rho, silva_deepening_history in silva_deepening_histories.items():
silva_deepening_axes[0].semilogy(
range(1, len(silva_deepening_history) + 1),
silva_deepening_history,
marker="o",
markersize=2,
linewidth=1.2,
label=f"rho={silva_deepening_rho:.2f}",
)
silva_deepening_axes[0].axhline(
silva_deepening_tolerance, color="black", linestyle="--", linewidth=0.9
)
silva_deepening_axes[0].set_xlabel("iteration")
silva_deepening_axes[0].set_ylabel("absolute residual")
silva_deepening_axes[0].set_title("Residual trajectories")
silva_deepening_axes[0].legend(fontsize=7)
silva_deepening_axes[1].plot(
[row[0] for row in silva_deepening_rows],
[row[1] for row in silva_deepening_rows],
marker="o",
label="iterations",
)
silva_deepening_sensitivity_axis = silva_deepening_axes[1].twinx()
silva_deepening_sensitivity_axis.plot(
[row[0] for row in silva_deepening_rows],
[row[4] for row in silva_deepening_rows],
color="tab:red",
marker="s",
label="sensitivity",
)
silva_deepening_axes[1].set_xlabel('implicit dynamics factor')
silva_deepening_axes[1].set_ylabel("iterations")
silva_deepening_sensitivity_axis.set_ylabel("implicit sensitivity", color="tab:red")
silva_deepening_axes[1].set_title("Cost and sensitivity")
silva_deepening_figure.tight_layout()
silva_deepening_plt.show()
Reading and Extending the Result¶
The measured residual curves become flatter as the implicit dynamics factor increases. The iteration count and the exact sensitivity rise together, but they answer different questions: iterations measure numerical work, while sensitivity describes how strongly the equilibrium reacts to the source. The gradient-error column verifies the differentiation path against the analytic derivative.
Apply the same separation to this notebook's full model:
| Report | Notebook-specific interpretation |
|---|---|
| Task evidence | equation residual, boundary error, trajectory error, and solver residual |
| Forward residual | Re-evaluate the complete transition at the returned state |
| Empirical rate | Compare consecutive residuals only after the transient regime |
| Backward residual | Record the linear-adjoint stopping value independently |
| Sensitivity | Perturb one declared source field while preserving all other inputs |
| Structural checks | initial/boundary conditions, equation shape, and constraint consistency |
| Scale sweep | Change one of collocation count, latent dimension, stages, stiffness, and time horizon at a time |
A richer experiment should now repeat the sweep with at least two forward solvers, two tolerances, and multiple seeds. Keep model parameters and data identical when comparing solvers. Then change one architecture or data-scale axis, retain the compact analytic study as a regression test, and report task quality, residuals, iterations, runtime, memory, gradient norms, and failed convergence cases together.