Stacking, Solvers, and Devices

SILVA layers are ordinary torch.nn.Module objects. A model can contain one equilibrium layer, a stack of equilibrium layers, a prediction head, or custom operators inside each equilibrium block.

Stacked Equilibrium Blocks

from silva_networks import SILVAStack, SolverConfig

stack = SILVAStack(
    in_dim=6,
    hidden_dims=[32, 32, 16],
    config=[
        SolverConfig(solver="picard", max_iter=12, alpha=0.5),
        SolverConfig(solver="anderson", max_iter=12, alpha=0.5, history=4),
        SolverConfig(solver="broyden", max_iter=8, alpha=0.4),
    ],
    local=["graph", "topk", "graph"],
    global_term="mean",
)

The first layer solves for a 32-channel equilibrium state. The second layer uses that state as its stimulus and solves another 32-channel equilibrium. The third layer maps the equilibrium representation to 16 channels. Each layer has its own SolverConfig, so a stable Picard block, an accelerated Anderson block, and a small Broyden block can live in the same architecture. Each config can also choose backward_mode="unrolled" or backward_mode="implicit".

For a stack of \(m\) points, the equations are

\[ z_1^\star=f_{\theta_1}(z_1^\star,x), \qquad z_\ell^\star=f_{\theta_\ell}(z_\ell^\star,z_{\ell-1}^\star), \quad 2\le\ell\le m. \]

These are \(m\) separate fixed points connected by explicit links. Adding modules to state_network instead adds depth inside one repeated transition; it does not create more equilibrium points.

For SILVA cortex hierarchies where a single equilibrium point also contains a deep internal transition network, use Cortex Hierarchies. That API keeps the same solver and device contract while allowing each point to have its own encoder, internal modules, interaction terms, link function, and alpha.

End-to-End Graph Models

from silva_networks import SILVAGraphNetwork, SolverConfig

model = SILVAGraphNetwork(
    in_dim=8,
    hidden_dims=[64, 64],
    out_dim=5,
    task="graph",
    pooling="mean",
    config=SolverConfig(solver="anderson", max_iter=20, alpha=0.5),
    local="graph",
    global_term="mean",
    head_hidden_dims=(64,),
)

For node prediction, use task="node". For graph or set prediction, use task="graph" and provide a batch vector when several graphs are packed into one tensor.

Custom Operators

Any trainable module can become the local or global branch if it returns a tensor with the same shape as the equilibrium state.

import torch

class SignedLocal(torch.nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.proj = torch.nn.Linear(dim, dim, bias=False)

    def forward(self, z, edge_index=None, batch=None, x=None):
        messages = self.proj(z)
        if edge_index is None:
            return messages
        src, dst = edge_index
        out = torch.zeros_like(messages)
        out.index_add_(0, dst, messages[src])
        return out

model = SILVAGraphNetwork(
    in_dim=8,
    hidden_dims=[32, 32],
    out_dim=3,
    local=lambda dim, layer_index: SignedLocal(dim),
    global_term="mean",
)

The SILVA wrapper passes available context by keyword: z, x, edge_index, and batch. A custom module may accept only the arguments it needs.

Device Execution

from silva_networks import move_to_device, resolve_device

device = resolve_device("auto")
model = model.to(device)
batch = move_to_device(batch, device)

Every internal zero tensor, identity matrix, aggregation buffer, and solver workspace follows the device and dtype of the input state. The same model code runs on CPU, CUDA, or MPS when the corresponding PyTorch backend is available.

Inspect Every Point

result = model(x, edge_index=edge_index, batch=batch, return_result=True)

for index, solve in enumerate(result.solver_results):
    print(index, solve.solver, solve.iterations, solve.converged, solve.residual)

Validate each point independently. A small final-layer residual does not prove that an earlier state converged, and a device check should include features, edges, batch ids, targets, model parameters, and every returned state.

The runnable construction is in Stacked Architecture. Solver sources are collected in Solvers and Linear Algebra.

Worked Evidence Bridge

The derivation above becomes a complete SILVA study when the state, condition, solver result, task result, and gradient path are kept separate. Here the state is three linked equilibrium states with independent solver policies and the condition is the output of the preceding point plus the original task input. The compact relation is

\[ z_i^\star=T_{\theta_i}(z_i^\star;c_i),\qquad c_{i+1}=A_i(z_i^\star,x) \]

The following is the complete executable program used by the repository tests:

from __future__ import annotations

import torch

from silva_networks import SILVAGraphNetwork, SolverConfig, move_to_device, resolve_device


class SignedLocal(torch.nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.proj = torch.nn.Linear(dim, dim, bias=False)

    def forward(self, z: torch.Tensor, edge_index: torch.Tensor | None = None) -> torch.Tensor:
        messages = self.proj(z)
        if edge_index is None:
            return messages
        src, dst = edge_index
        out = torch.zeros_like(messages)
        out.index_add_(0, dst, messages[src])
        return out


def main() -> None:
    torch.manual_seed(19)
    device = resolve_device("auto")

    batch = {
        "x": torch.randn(14, 6),
        "edge_index": torch.tensor(
            [list(range(13)), list(range(1, 14))],
            dtype=torch.long,
        ),
        "batch": torch.tensor([0] * 7 + [1] * 7),
        "y": torch.tensor([0, 1]),
    }
    batch = move_to_device(batch, device)

    model = SILVAGraphNetwork(
        in_dim=6,
        hidden_dims=[16, 16, 12],
        out_dim=2,
        task="graph",
        pooling="mean",
        config=[
            SolverConfig(solver="picard", max_iter=8, alpha=0.5),
            SolverConfig(solver="anderson", max_iter=8, alpha=0.5, history=3),
            SolverConfig(solver="broyden", max_iter=8, alpha=0.5),
        ],
        local=lambda dim, index: SignedLocal(dim) if index == 1 else "graph",
        global_term="mean",
        head_hidden_dims=(16,),
    ).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
    for _ in range(3):
        result = model(
            batch["x"],
            edge_index=batch["edge_index"],
            batch=batch["batch"],
            return_results=True,
        )
        loss = torch.nn.functional.cross_entropy(result.output, batch["y"])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    solvers = [solver_result.solver for solver_result in result.solver_results or []]
    print("device", device.type)
    print("logits_shape", tuple(result.output.shape))
    print("solvers", solvers)
    print("final_loss", float(loss.detach().cpu()))


if __name__ == "__main__":
    main()

Run it from the project root:

python examples/stacked_architecture.py

Measured Output

device cpu
logits_shape (2, 2)
solvers ['picard', 'anderson', 'broyden']
final_loss 0.586849570274353

What This Result Establishes

This run records the final logit shape, all three selected solvers, and a differentiable task loss. It establishes that the compact mechanism is executable with finite outputs and that its stated shape or structural contract can be inspected. It does not establish source-scale accuracy by itself.

For the next controlled study, profile each point separately before increasing point count, state width, batch size, or device count. Keep the compact run as a regression case. For every larger run, archive the resolved data source and split, preprocessing, seed, constructor arguments, forward and backward solver settings, task metric, normalized residual, iteration count, gradient norm, runtime, peak memory, and convergence failures. This keeps task quality, numerical convergence, and computational cost from being collapsed into one number.

Where to Go Next

Question	Page
How are linked points organized as a cortex hierarchy?	Cortex Hierarchies
Where is a stacked model executed?	Stacked Architecture Example
Which modules and device helpers are public?	Architectures API