Unrolling a cycle¶
This page goes through every computation that runs a cyclic graph in PyTorch, one step at a time. It uses the five-node graph of Cyclic graph example with neutral input names, sets every weight by hand, and draws each step, so you can check every number with pen and paper.
Read it when you want to see how a feedback edge runs, how many
steps a value needs to reach the output, and what decides when the
loop stops. To train and interpret a model on a cyclic graph, use
the Cyclic graph example instead. A
directed acyclic graph (DAG) needs no loop:
parse_layered() computes it in one sweep, see
Layered vs. Adjacency.
The page shows five facts:
- Every step updates all nodes at once. Each node reads the values its source nodes had after the previous step.
- A value moves one edge per step.
- A node keeps nothing of its own previous value unless a cycle or
a self-loop brings it back.
outputis overwritten every step. - After $K$ steps,
outputholds every path from an input that is at most $K$ edges long. - The cycle runs across steps, never inside one step, so autograd sees an ordinary feedforward computation.
kpnn2 builds the spec — the parsed layout
of the graph — and the MaskedLinear. The loop is your own code,
and so is the choice of when it stops (see
When to stop).
import random
import numpy as np
import pandas as pd
import torch
from graphviz import Digraph
from matplotlib import pyplot as plt
from matplotlib_inline.backend_inline import (
set_matplotlib_formats,
)
import kpnn2
torch.manual_seed(42)
random.seed(42)
np.random.seed(42)
set_matplotlib_formats(
"svg",
bbox_inches=None,
)
plt.rcParams.update(
{
"font.size": 14,
"font.family": "Liberation Sans",
"lines.linewidth": 1.5,
"axes.facecolor": "none",
"figure.facecolor": "none",
"savefig.facecolor": "none",
"svg.fonttype": "none",
}
)
edgelist = pd.DataFrame(
{
"source": [
"input_a",
"input_b",
"node_a",
"node_b",
"node_a",
],
"target": [
"node_a",
"node_b",
"node_b",
"node_a",
"output",
],
}
)
spec = kpnn2.parse_adjacency(edgelist)
spec.nodes, spec.input_index, spec.output_index
(('input_a', 'input_b', 'node_a', 'node_b', 'output'), (0, 1), (4,))
parse_adjacency() puts every node into one
state vector — one slot per node — in
the alphabetical order of spec.nodes. The cycle needs nothing
special: the edge node_b → node_a is one more pair of indices,
like every other edge. spec.to_mask() holds these edges as a
square, one row per node and a 1 for each node it reads. The input
rows are empty, because no edge enters an input. So the loop has to
write the input values in before every step.
Weights you can check by hand¶
Set every weight to 1 and every bias to 0, and feed input_a = 10
and input_b = 1. Each path then adds its input value unchanged.
The mask keeps absent edges at zero, so filling the whole weight
square with 1 is safe.
core = kpnn2.MaskedLinear(spec.to_mask())
with torch.no_grad():
core.weight.fill_(1.0)
core.bias.zero_()
x_df = pd.DataFrame(
{
"input_a": [10.0],
"input_b": [1.0],
}
)
col = kpnn2.align_inputs(
x_df.columns,
spec,
)
x = torch.as_tensor(
x_df.to_numpy()[:, col],
dtype=torch.float32,
)
x
tensor([[10., 1.]])
The loop is the forward() of CyclicNet in the Cyclic graph
example, without the head: write the inputs in, apply core,
repeat. unroll() keeps the state after every step, with the zero
start as step 0, so states has the dims (step, sample, node).
def unroll(
spec,
core,
x,
n_steps,
):
input_index = torch.as_tensor(spec.input_index)
state = x.new_zeros(
x.shape[0],
core.in_features,
)
states = [state]
for _ in range(n_steps):
state = state.index_copy(
-1,
input_index,
x,
)
state = core(state)
states.append(state)
return torch.stack(states)
with torch.no_grad():
states = unroll(
spec,
core,
x,
n_steps=4,
)
states.shape
torch.Size([5, 1, 5])
The loop unrolled over steps¶
Draw every step as its own column, and the cycle turns into a
pattern that repeats between columns. Each arrow runs from a node in
one column to a node in the next column, one arrow per edge, because
a step reads only the state the previous step left. The feedback
pair node_a ⇄ node_b is the crossing between every two columns.
The numbers come from states. Dashed boxes are the input values
the loop writes in before the next step.
Pick any box and add up the boxes its arrows come from: the sum is
the number in the box. node_a after step 2 reads 10 from input_a
and 1 from node_b, and holds 11. With other weights, each source
counts weight × value, and the bias is added on top. In these
figures the tens digit counts the paths from input_a and the ones
digit the paths from input_b.
Follow one value through it. input_a = 10 enters node_a in step
1, and output reads that 10 in step 2. input_b = 1 enters
node_b in step 1, reaches node_a in step 2 and output in step
3, where 11 = 10 + 1. output is still 0 after step 1, because in
step 1 it reads node_a from the start.
No arrow points back to an earlier column or stays inside one, so the computation itself has no cycle.
Which paths reach the output¶
output after step 4 depends only on the boxes from which a chain
of arrows leads to it. The figure below draws exactly those arrows
dark and thick.
Three paths run along the dark arrows:
input_a → node_a → outputadds 10.input_b → node_b → node_a → outputadds 1.input_a → node_a → node_b → node_a → outputgoes once around the loop and adds 10.
Together they give 21, the value of output after step 4. Each
extra step adds the paths that are one edge longer. The dark arrow
from node_b at the start carries 0. The inputs written before the
last step have no dark arrow: they reach only the nodes one edge
away, never output.
With weights other than 1, each path adds its input value times the product of the weights along it. A nonlinear activation breaks this clean sum, but not which paths can contribute.
What autograd sees¶
Autograd records the computation the way the figure draws it: one
application of core per step, each reading the state of the step
before. Backpropagation follows the dark arrows backwards, so an
input gets a gradient only if a dark arrow leaves it. Here are the
dark arrows after two steps:
No dark arrow leaves input_b. With two steps, input_b gets no
gradient, so training cannot learn anything from it. With three or
more steps a dark arrow leaves input_b, as in the step-4 figure
above. This is why the controls in the Cyclic graph example run
three steps: the path from the input of node_b to output is
three edges long.
Without the feedback edge¶
Remove the edge node_a → node_b, and the graph is a DAG. Its
longest path from an input to output,
input_b → node_b → node_a → output, has three edges. Run the same
loop on it with the same hand-set weights:
is_feedback = (edgelist["source"] == "node_a") & (
edgelist["target"] == "node_b"
)
dag_edgelist = edgelist[~is_feedback]
dag_spec = kpnn2.parse_adjacency(dag_edgelist)
dag_core = kpnn2.MaskedLinear(dag_spec.to_mask())
with torch.no_grad():
dag_core.weight.fill_(1.0)
dag_core.bias.zero_()
# Same input nodes as spec, so x still lines up.
dag_states = unroll(
dag_spec,
dag_core,
x,
n_steps=5,
)
From step 3 on, every column is the same: no path is longer than
three edges, so every path has arrived. On a DAG the loop settles by
itself once the step count reaches the longest path, and more steps
cost time but change nothing. For a DAG, parse_layered() with one
sweep through its layers is the direct route. A cycle makes the
paths endless, so with a cycle something has to decide when the loop
stops.
The loop in formulas¶
Write $n$ for the number of nodes, 5 here. For one sample:
- $h_t \in \mathbb{R}^n$ is the state after step $t$ (
states[t, 0]above), with $h_0 = 0$. - $W \in \mathbb{R}^{n \times n}$ is the weight that
coreapplies,core.effective_weight(). $W_{ij}$ is the weight of the edge $j \to i$, and 0 where there is no edge. - $b \in \mathbb{R}^n$ is the bias.
- $\hat{x} \in \mathbb{R}^n$ holds the input values in the input slots and 0 elsewhere.
One step overwrites the input slots with $\hat{x}$ and applies
core. What the input slots held before is lost, so it never
reaches $W$. Let $A$ be $W$ with its input columns set to 0, and let
$u = W \hat{x} + b$. Then
$$ \begin{aligned} h_t &= A \, h_{t-1} + u, \\ h_K &= u + A u + A^2 u + \cdots + A^{K-1} u = \sum_{k=0}^{K-1} A^k u. \end{aligned} $$
$u$ holds the one-edge paths out of the inputs, and every product with $A$ makes each path one edge longer. So $h_K$ holds every path of at most $K$ edges, as the figures showed. If the spectral radius $\rho(A)$, the largest magnitude among the eigenvalues of $A$, is below 1, the sum converges as $K$ grows, to the steady state
$$ h^* = (I - A)^{-1} u, $$
where $I$ is the $n \times n$ identity. With an activation $f$, a step is $h_t = f(A \, h_{t-1} + u)$. The sum no longer splits into paths, but the steady state is still the state that one more step leaves unchanged, $h^* = f(A \, h^* + u)$.
loop_terms() builds $A$ and $u$ from a MaskedLinear:
def loop_terms(
spec,
core,
x,
):
weight = core.effective_weight().detach()
x_hat = torch.zeros(core.in_features)
x_hat[list(spec.input_index)] = x[0]
a = weight.clone()
a[:, list(spec.input_index)] = 0.0
u = weight @ x_hat + core.bias.detach()
return a, u
a, u = loop_terms(
spec,
core,
x,
)
torch.linalg.eigvals(a).abs().max()
tensor(1.)
With every weight 1, the loop node_a ⇄ node_b gives $A$ the
eigenvalues 1 and −1, so $\rho(A) = 1$. The sum never
converges, and the values in the figures keep growing.
When to stop¶
Everything above follows from the graph. What differs between
models on a cyclic graph is when the loop stops. kpnn2 ships none
of these loops: the loop, the stopping rule, any convergence check
or stability penalty, and any custom backward pass are your own
code. Three choices are common.
A fixed number of steps¶
The Cyclic graph example and this page use a fixed number of steps
$K$. output holds every path of at most $K$ edges, and autograd
backpropagates through $K$ applications of core. Pick $K$ at least
as large as the number of edges on the longest input-to-output path
the model should use. A shorter $K$ cuts those inputs off, as the
two-step figure shows. The result depends on $K$, and when the loop
amplifies, as here, values keep growing with $K$.
Run until the state stops changing¶
Repeat the step until the change from one step to the next falls below a tolerance, then read the steady state $h^*$. The result no longer depends on a step count. For a signaling network it reads as the state the network settles in under a constant input. Without an activation, the loop reaches $h^*$ only if $\rho(A) < 1$, that is, only if the loop shrinks what goes around it.
Set both feedback weights to 0.5 instead of 1. A lap around
node_a ⇄ node_b then multiplies by 0.25, and the loop settles.
run_to_steady_state() stops once every slot changes by less than
1e-6, and returns None as the step count if that never happens
within max_steps:
def run_to_steady_state(
spec,
core,
x,
tolerance=1e-6,
max_steps=500,
):
input_index = torch.as_tensor(spec.input_index)
state = x.new_zeros(
x.shape[0],
core.in_features,
)
for step in range(1, max_steps + 1):
previous = state
state = core(
state.index_copy(
-1,
input_index,
x,
)
)
if (state - previous).abs().max() < tolerance:
return state, step
return state, None
damped = kpnn2.MaskedLinear(spec.to_mask())
with torch.no_grad():
damped.weight.fill_(1.0)
damped.bias.zero_()
damped.weight[
spec.node_units("node_b"),
spec.node_units("node_a"),
] = 0.5
damped.weight[
spec.node_units("node_a"),
spec.node_units("node_b"),
] = 0.5
steps_to_settle = {
name: run_to_steady_state(
spec,
layer,
x,
)[1]
for name, layer in [
("feedback weights 1", core),
("feedback weights 0.5", damped),
]
}
steps_to_settle
{'feedback weights 1': None, 'feedback weights 0.5': 25}
With feedback weights 1 the loop never settles. With feedback weights 0.5 it settles after 25 steps. The plot at the end of this section shows where.
Backpropagating through every step of a long run stores every step.
A steady-state model can instead compute the gradient at $h^*$ with an
implicit (adjoint) backward pass, a custom torch.autograd.Function
you write. To keep the loop stable during training, penalize the
spectral radius in your loss or constrain the weights.
LEMBAS — large-scale knowledge-embedded artificial signaling networks (Nilsson et al., 2022) — is built this way on a signaling network. It has one state per node and one weight per edge. All nodes update together from zero, with the ligand input added at every step, until the state stops changing or a step limit is reached. It scales its weights to a spectral radius below 1 at the start, penalizes the spectral radius during training, and computes gradients at the steady state with its own backward pass instead of storing every step.
Integrate in continuous time¶
Each step of the loop is one Euler step of size 1 for an ordinary differential equation (ODE) in a time $\tau$:
$$ \frac{\mathrm{d}h}{\mathrm{d}\tau} = A \, h + u - h. $$
An Euler step of size $\Delta\tau$ moves $h$ to $h + \Delta\tau \, (A \, h + u - h)$, and $\Delta\tau = 1$ gives $A \, h + u$, the loop. A smaller step, or an adaptive solver, follows the same equation more finely. With an activation $f$, the right-hand side is $f(A \, h + u) - h$. The equation comes to rest where $\mathrm{d}h / \mathrm{d}\tau = 0$, at the same steady state $h^*$. Without an activation it gets there if every eigenvalue of $A$ has real part below 1, which $\rho(A) < 1$ implies. The stopping rule becomes an end time or, again, a steady state.
To use it, write the right-hand side as a module that calls
core, and integrate it with your own Euler loop or with an ODE
solver such as torchdiffeq.
integrate_euler() is the Euler loop, and it keeps the state at
every time point:
def integrate_euler(
spec,
core,
x,
step_size,
end_time,
):
input_index = torch.as_tensor(spec.input_index)
state = x.new_zeros(
x.shape[0],
core.in_features,
)
states = [state]
for _ in range(round(end_time / step_size)):
state = state.index_copy(
-1,
input_index,
x,
)
state = state + step_size * (core(state) - state)
states.append(state)
return torch.stack(states)
The three rules side by side¶
Follow output for 30 steps under each rule: the loop with feedback
weights 1, the loop with feedback weights 0.5, and the Euler loop
with step size 0.1 on the feedback-0.5 weights. The second legend
lists reference marks, not runs: the steady state from the formula,
and the step at which a loop with a fixed $K = 4$ reads output.
n_steps = 30
step_size = 0.1
with torch.no_grad():
growing = unroll(
spec,
core,
x,
n_steps,
)
settling = unroll(
spec,
damped,
x,
n_steps,
)
euler = integrate_euler(
spec,
damped,
x,
step_size,
end_time=float(n_steps),
)
a, u = loop_terms(
spec,
damped,
x,
)
steady_state = torch.linalg.solve(
torch.eye(len(u)) - a,
u,
)
The loop with feedback weights 1 leaves the plot after a few steps and keeps growing. With feedback weights 0.5, the loop and the Euler loop both come to rest on the dashed line, the steady state the formula gives. The rings mark what a loop with a fixed $K = 4$ reads: 21 with feedback weights 1, and 13 with feedback weights 0.5, short of the steady state. The Euler curve has no ring, because its x-axis is time: reaching $\tau = 4$ takes 40 Euler steps of size 0.1.
The three rules agree only when the loop settles and $K$ is large enough to come close to where it settles. A fixed $K$ reads the sum after $K$ terms, wherever it stands. The steady-state and ODE rules read where the loop comes to rest, and they require that it does.
Where to go next¶
- Cyclic graph example trains this loop on simulated data and scores the nodes at every step.
- Time-series example writes a different
x_tin at each step instead of the samex. - Layered vs. Adjacency compares this
layout with
parse_layered(). - PackedLinear replaces
MaskedLinearwhen the(n, n)square would strain RAM.
References¶
- Nilsson A, Peters JM, Meimetis N, Bryson B, Lauffenburger DA (2022). Artificial neural networks enable genome-scale simulations of intracellular signaling. Nature Communications 13:3069. doi:10.1038/s41467-022-30684-y