SILVA Distributional Equilibrium: Empirical-Measure Lab¶
This lab derives measure discrepancies, verifies permutation behavior and
masks, runs particle descent, and trains a small task readout from a
distributional SILVA state. The canonical family is
silva_distributional_deq [45].
Numbered literature: [1], [45]. Each number opens the complete citation and its primary external source.
from pathlib import Path
import importlib.util
import subprocess
import sys
IN_HOSTED_RUNTIME = "google.colab" in sys.modules
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 IN_HOSTED_RUNTIME and 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
from torch import nn
from silva_networks import (
SILVADistributionalDEQ,
distributional_discrepancy,
make_variable_measure_dataset,
silva_equilibrium_model,
)
plt.rcParams.update({"figure.dpi": 300, "savefig.dpi": 300})
torch.manual_seed(200)
<torch._C.Generator at 0x114d04eb0>
1. Matrices Represent Empirical Measures¶
Input rows $X=(x_1,\ldots,x_M)$ represent
$$ \rho_X=\frac1M\sum_{i=1}^{M}\delta_{x_i}. $$
Latent rows $Z=(z_1,\ldots,z_N)$ represent
$$ \mu_Z=\frac1N\sum_{j=1}^{N}\delta_{z_j}. $$
Permuting rows does not change either measure. Padded storage therefore needs a boolean mask so padding contributes neither to attention nor to empirical expectations.
data = make_variable_measure_dataset(
samples=6,
min_particles=5,
max_particles=9,
dimension=2,
components=2,
seed=20,
)
assert torch.equal(data.context_mask.sum(dim=1), data.counts)
assert torch.allclose(data.empirical_mean(), data.target_mean)
print("context:", tuple(data.context.shape))
print("particle counts:", data.counts.tolist())
print("target means:", data.target_mean)
context: (6, 9, 2)
particle counts: [6, 5, 7, 8, 9, 5]
target means: tensor([[ 0.6399, -1.3564],
[-0.0019, 0.8456],
[-0.6017, 0.1548],
[ 0.8623, -0.2763],
[-0.7654, 1.3412],
[-0.3914, -0.9826]])
2. Gaussian MMD and Energy Distance¶
For kernel $k$, the biased squared maximum mean discrepancy is
$$ \begin{aligned} \operatorname{MMD}^2(\mu,\nu) &=\mathbb E_{x,x'\sim\mu}k(x,x')\\ &\quad+\mathbb E_{y,y'\sim\nu}k(y,y')\\ &\quad-2\mathbb E_{x\sim\mu,y\sim\nu}k(x,y). \end{aligned} $$
SILVA provides the Gaussian kernel
$$ k_\ell(x,y)=\exp\left(-\frac{\|x-y\|^2}{2\ell^2}\right) $$
and the energy distance
$$ D_E^2 =2\mathbb E\|x-y\| -\mathbb E\|x-x'\| -\mathbb E\|y-y'\|. $$
permutation = torch.tensor([4, 0, 3, 1, 2, 5, 6, 7, 8])
original = distributional_discrepancy(
data.context,
data.context,
kernel="gaussian",
left_mask=data.context_mask,
right_mask=data.context_mask,
)
permuted = distributional_discrepancy(
data.context,
data.context[:, permutation],
kernel="gaussian",
left_mask=data.context_mask,
right_mask=data.context_mask[:, permutation],
)
assert torch.allclose(original, permuted, atol=1e-6)
print("self discrepancy:", float(original))
print("permuted discrepancy:", float(permuted))
self discrepancy: 0.0 permuted discrepancy: 0.0
3. Distributional SILVA Objective¶
The transition maps a latent measure and input measure to a transformed latent measure. The equilibrium objective is
$$ G_{\theta,X}(Z) =\frac12D^2\left(\mu_Z,\mu_{F_\theta(Z,X)}\right). $$
Particle descent applies
$$ z_j^{k+1} =z_j^k-\eta\nabla_{z_j}G_{\theta,X}(Z^k). $$
For latent permutation $P$ and context permutation $Q$, the built-in transition satisfies
$$ F_\theta(PZ,QX)=P F_\theta(Z,X). $$
This is equivariance in latent order and invariance in context order [45].
model = SILVADistributionalDEQ(
input_dim=2,
latent_dim=4,
particles=5,
heads=2,
kernel="gaussian",
bandwidth=1.0,
step_size=0.12,
max_iter=3,
)
result = model(
data.context[:2],
context_mask=data.context_mask[:2],
return_result=True,
)
assert result.state.shape == (2, 5, 4)
assert result.discrepancies[-1] <= result.discrepancies[0] + 1e-6
print("discrepancy path:", result.discrepancies)
print("converged to configured tolerance:", result.converged)
discrepancy path: [0.5253183245658875, 0.522333025932312, 0.49297821521759033, 0.48192644119262695] converged to configured tolerance: False
4. Train a Readout from the Equilibrium Measure¶
The equilibrium state remains a set. A task head must pool in a way consistent with the target. Here a pointwise decoder maps latent particles back to two dimensions, and their mean predicts the empirical input mean:
$$ \widehat m =\frac1N\sum_{j=1}^{N}Q_\psi(z_j^\star). $$
decoder = nn.Linear(4, 2)
optimizer = torch.optim.Adam(
[*model.parameters(), *decoder.parameters()],
lr=3e-3,
)
losses = []
for epoch in range(4):
optimizer.zero_grad()
state = model(data.context, context_mask=data.context_mask)
decoded_particles = decoder(state)
prediction = decoded_particles.mean(dim=1)
task_loss = torch.nn.functional.mse_loss(prediction, data.target_mean)
transformed = model.transition(
state,
data.context,
context_mask=data.context_mask,
)
equilibrium_loss = distributional_discrepancy(
state,
transformed,
kernel="gaussian",
)
loss = task_loss + 0.05 * equilibrium_loss
loss.backward()
optimizer.step()
losses.append(float(loss.detach()))
assert all(torch.isfinite(torch.tensor(losses)))
print("training losses:", losses)
print("task loss:", float(task_loss.detach()))
print("distributional discrepancy:", float(equilibrium_loss.detach()))
training losses: [0.8673074245452881, 0.8647595047950745, 0.8614999055862427, 0.859478235244751] task loss: 0.8103713989257812 distributional discrepancy: 0.9821367859840393
with torch.no_grad():
final_state = model(data.context[:1], context_mask=data.context_mask[:1])
decoded = decoder(final_state)[0]
valid_context = data.context[0, data.context_mask[0]]
fig, axes = plt.subplots(1, 2, figsize=(6.2, 2.7))
axes[0].scatter(valid_context[:, 0], valid_context[:, 1], label="context")
axes[0].scatter(decoded[:, 0], decoded[:, 1], marker="x", label="latent readout")
axes[0].legend()
axes[0].set_title("empirical measures")
axes[1].plot(range(1, len(losses) + 1), losses, marker="o")
axes[1].set_yscale("log")
axes[1].set_xlabel("epoch")
axes[1].set_title("training objective")
fig.tight_layout()
plt.show()
5. Fixed and Invalid Particles¶
latent_mask identifies valid latent rows. fixed_mask is a subset that
remains equal to its initial value during particle descent. This supports
observed anchors or boundary particles without allowing padding rows to enter
the measure.
z0 = torch.randn(1, 5, 4)
latent_mask = torch.tensor([[True, True, True, True, False]])
fixed_mask = torch.tensor([[True, False, False, False, False]])
anchored = model(
data.context[:1],
z0=z0,
context_mask=data.context_mask[:1],
latent_mask=latent_mask,
fixed_mask=fixed_mask,
return_result=True,
)
assert torch.equal(anchored.state[:, 0], z0[:, 0])
assert torch.equal(anchored.state[:, 4], torch.zeros_like(anchored.state[:, 4]))
print("fixed particle preserved and padding excluded")
fixed particle preserved and padding excluded
factory_model = silva_equilibrium_model(
"silva_distributional_deq",
input_dim=2,
latent_dim=4,
particles=4,
heads=2,
max_iter=1,
)
factory_state = factory_model(
data.context[:1],
context_mask=data.context_mask[:1],
)
assert factory_state.shape == (1, 4, 4)
print(type(factory_model).__name__)
SILVADistributionalDEQ
6. Practical Guidance¶
| Problem | Diagnostic | Response |
|---|---|---|
| result changes after row permutation | run the EI check | remove positional row encodings and order-dependent pooling |
| padding changes the result | compare variable-length and padded forms | pass masks through attention and discrepancy terms |
| discrepancy oscillates | inspect every descent step | reduce particle step size or change bandwidth |
| task loss is low but measures disagree | report task and measure losses separately | retain an explicit equilibrium diagnostic |
The generated mixtures validate variable-size storage, masks, invariance, and outer gradients. Reproducing point-cloud benchmarks requires their official splits, augmentations, sample counts, and task metrics [45].
From 20 Silva Distributional Equilibrium Lab 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 | a masked set of latent particles |
| Condition | an empirical input measure and validity mask |
| Repeated computation | a permutation-compatible particle transition and discrepancy descent |
| Required invariants | permutation equivariance, masks, variable cardinality, and finite particles |
| Replaceable components | particle initializer, attention transition, discrepancy, descent rule, and readout |
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.
Replace the Particle Transition¶
class MyMeasureTransition(nn.Module):
def forward(self, latent, context, *, latent_mask=None, context_mask=None):
return self.permutation_compatible_update(
latent, context, latent_mask, context_mask
)
model = SILVADistributionalDEQ(
input_dim=input_dim,
latent_dim=latent_dim,
transition=MyMeasureTransition(),
kernel="energy",
pairwise_chunk_size=pairwise_chunk_size,
)
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 measure discrepancy, moment error, task error, and descent 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 particle count, latent width, pair chunk size, and batch cardinality. 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": '20_silva_distributional_equilibrium_lab.ipynb',
"state": 'a masked set of latent particles',
"condition": 'an empirical input measure and validity mask',
"transition": 'a permutation-compatible particle transition and discrepancy descent',
"invariants": 'permutation equivariance, masks, variable cardinality, and finite particles',
"compact_metric": 'measure discrepancy, moment error, task error, and descent residual',
"scale_axis": 'particle count, latent width, pair chunk size, and batch cardinality',
}
assert all(notebook_reproduction_record.values())
notebook_reproduction_record
{'notebook': '20_silva_distributional_equilibrium_lab.ipynb',
'state': 'a masked set of latent particles',
'condition': 'an empirical input measure and validity mask',
'transition': 'a permutation-compatible particle transition and discrepancy descent',
'invariants': 'permutation equivariance, masks, variable cardinality, and finite particles',
'compact_metric': 'measure discrepancy, moment error, task error, and descent residual',
'scale_axis': 'particle count, latent width, pair chunk size, and batch cardinality'}
Worked Convergence and Sensitivity Study¶
The preceding example demonstrates one configured solve. This additional study changes the measure-coupling 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, a masked set of latent particles, and its repeated map, a permutation-compatible particle transition and discrepancy descent. 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('measure-coupling 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)
measure-coupling 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('measure-coupling 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 measure-coupling 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 | measure discrepancy, moment error, task error, and descent 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 | permutation equivariance, masks, variable cardinality, and finite particles |
| Scale sweep | Change one of particle count, latent width, pair chunk size, and batch cardinality 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.