State-update backend example¶

This notebook walks through an end-to-end workflow with the state_update backend: compile a cyclic graph, train a small model, and interpret named nodes.

It assumes you have already read the Getting started notebook. That tutorial covers package basics and the overall edge2torch workflow. Here we focus on what changes when the graph contains cycles and the model is compiled as a fixed-step state-update network.

In [1]:

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import torch
from graphviz import Digraph
from IPython.display import display
from torch import nn

import edge2torch as e2t

sns.set_theme(style="whitegrid")

import matplotlib.pyplot as plt import numpy as np import pandas as pd import seaborn as sns import torch from graphviz import Digraph from IPython.display import display from torch import nn import edge2torch as e2t sns.set_theme(style="whitegrid")

Define a cyclic graph¶

The state_update backend allows cycles among internal nodes. In this toy example, two hidden nodes exchange feedback while input features feed each node:

• input_signal is the input to node_a
• input_noise is the input to node_b
• node_a and node_b form a feedback loop
• node_a connects to the output node output

The first walkthrough below uses only input_signal as informative. A later section repeats the same architecture under positive and negative controls that vary which inputs carry signal.

In [2]:

edgelist = pd.DataFrame(
    {
        "source": [
            "input_signal",
            "input_noise",
            "node_a",
            "node_b",
            "node_a",
        ],
        "target": [
            "node_a",
            "node_b",
            "node_b",
            "node_a",
            "output",
        ],
    }
)

edgelist

edgelist = pd.DataFrame( { "source": [ "input_signal", "input_noise", "node_a", "node_b", "node_a", ], "target": [ "node_a", "node_b", "node_b", "node_a", "output", ], } ) edgelist

Out[2]:

	source	target
0	input_signal	node_a
1	input_noise	node_b
2	node_a	node_b
3	node_b	node_a
4	node_a	output

In [3]:

dot = Digraph()
dot.attr(rankdir="LR")

input_nodes = ["input_signal", "input_noise"]
state_nodes = ["node_a", "node_b"]
output_nodes = ["output"]

for node_name in input_nodes:
    dot.node(node_name, node_name)

for node_name in state_nodes:
    dot.node(node_name, node_name)

for node_name in output_nodes:
    dot.node(node_name, node_name)

for row in edgelist.itertuples(index=False):
    dot.edge(str(row.source), str(row.target))

dot

dot = Digraph() dot.attr(rankdir="LR") input_nodes = ["input_signal", "input_noise"] state_nodes = ["node_a", "node_b"] output_nodes = ["output"] for node_name in input_nodes: dot.node(node_name, node_name) for node_name in state_nodes: dot.node(node_name, node_name) for node_name in output_nodes: dot.node(node_name, node_name) for row in edgelist.itertuples(index=False): dot.edge(str(row.source), str(row.target)) dot

Out[3]:

Compile with backend="state_update"¶

Use the steps argument to control how many state-update steps are unrolled during compilation. Node interpretation sites are exposed as step_1, step_2, and so on.

In [4]:

model, artifact = e2t.compile_graph(
    edgelist=edgelist,
    backend="state_update",
    steps=2,
)

artifact.backend, list(artifact.interpretation_sites.keys())

model, artifact = e2t.compile_graph( edgelist=edgelist, backend="state_update", steps=2, ) artifact.backend, list(artifact.interpretation_sites.keys())

Out[4]:

('state_update', ['step_1', 'step_2'])

Simulate data, align features, and train¶

In [5]:

def simulate_binary_data(
    n_per_class,
    feature_names,
    informative_features,
    rng,
):
    n_features = len(feature_names)
    n_samples = 2 * n_per_class

    x = rng.normal(0.0, 1.0, size=(n_samples, n_features))
    y = np.array([0] * n_per_class + [1] * n_per_class)

    informative_idx = [
        feature_names.index(feature_name)
        for feature_name in informative_features
    ]

    for class_value, shift in enumerate([-1.5, 1.5]):
        class_mask = y == class_value
        for feature_idx in informative_idx:
            x[class_mask, feature_idx] += shift

    x_df = pd.DataFrame(x, columns=feature_names)
    y_series = pd.Series(y, name="label")

    return x_df, y_series

def simulate_binary_data( n_per_class, feature_names, informative_features, rng, ): n_features = len(feature_names) n_samples = 2 * n_per_class x = rng.normal(0.0, 1.0, size=(n_samples, n_features)) y = np.array([0] * n_per_class + [1] * n_per_class) informative_idx = [ feature_names.index(feature_name) for feature_name in informative_features ] for class_value, shift in enumerate([-1.5, 1.5]): class_mask = y == class_value for feature_idx in informative_idx: x[class_mask, feature_idx] += shift x_df = pd.DataFrame(x, columns=feature_names) y_series = pd.Series(y, name="label") return x_df, y_series

In [6]:

rng = np.random.default_rng(0)

x_train_df, y_train = simulate_binary_data(
    n_per_class=80,
    feature_names=artifact.feature_names,
    informative_features=["input_signal"],
    rng=rng,
)

x_test_df, y_test = simulate_binary_data(
    n_per_class=40,
    feature_names=artifact.feature_names,
    informative_features=["input_signal"],
    rng=rng,
)

x_train_df.head()

rng = np.random.default_rng(0) x_train_df, y_train = simulate_binary_data( n_per_class=80, feature_names=artifact.feature_names, informative_features=["input_signal"], rng=rng, ) x_test_df, y_test = simulate_binary_data( n_per_class=40, feature_names=artifact.feature_names, informative_features=["input_signal"], rng=rng, ) x_train_df.head()

Out[6]:

	input_noise	input_signal
0	0.125730	-1.632105
1	0.640423	-1.395100
2	-0.535669	-1.138405
3	1.304000	-0.552919
4	-0.703735	-2.765421

In [7]:

# Visualize the simulated training data

plot_df = x_train_df.copy()
plot_df["label"] = y_train.values
plot_df["class"] = plot_df["label"].map({0: "class 0", 1: "class 1"})

long_df = plot_df.melt(
    id_vars=["label", "class"],
    value_vars=artifact.feature_names,
    var_name="feature",
    value_name="value",
)

feature_type_map = {
    "input_signal": "informative",
    "input_noise": "uninformative",
}

long_df["feature_type"] = long_df["feature"].map(feature_type_map)

g = sns.catplot(
    data=long_df,
    x="class",
    y="value",
    col="feature",
    col_wrap=2,
    hue="class",
    kind="box",
    sharey=False,
    height=3.2,
    aspect=1.0,
    legend=False,
)

g.fig.suptitle(
    "Simulated training data: informative vs. uninformative features",
    y=0.98,
)

g.fig.subplots_adjust(top=0.78, hspace=0.35)

for ax, feature in zip(g.axes.flat, artifact.feature_names):
    feature_type = feature_type_map[feature]
    ax.set_title(f"{feature}\n({feature_type})")
    ax.set_xlabel("")
    ax.set_ylabel("value")

plt.show()

# Visualize the simulated training data plot_df = x_train_df.copy() plot_df["label"] = y_train.values plot_df["class"] = plot_df["label"].map({0: "class 0", 1: "class 1"}) long_df = plot_df.melt( id_vars=["label", "class"], value_vars=artifact.feature_names, var_name="feature", value_name="value", ) feature_type_map = { "input_signal": "informative", "input_noise": "uninformative", } long_df["feature_type"] = long_df["feature"].map(feature_type_map) g = sns.catplot( data=long_df, x="class", y="value", col="feature", col_wrap=2, hue="class", kind="box", sharey=False, height=3.2, aspect=1.0, legend=False, ) g.fig.suptitle( "Simulated training data: informative vs. uninformative features", y=0.98, ) g.fig.subplots_adjust(top=0.78, hspace=0.35) for ax, feature in zip(g.axes.flat, artifact.feature_names): feature_type = feature_type_map[feature] ax.set_title(f"{feature}\n({feature_type})") ax.set_xlabel("") ax.set_ylabel("value") plt.show()

In [8]:

torch.manual_seed(0)

x_train = e2t.align_features_to_input_nodes(
    data=x_train_df,
    artifact=artifact,
)
x_test = e2t.align_features_to_input_nodes(
    data=x_test_df,
    artifact=artifact,
)

customized_model = e2t.customize_model(
    model=model,
    head=nn.Linear(1, 1),
)

y_train_tensor = torch.tensor(
    y_train.values.reshape(-1, 1),
    dtype=torch.float32,
)

optimizer = torch.optim.Adam(customized_model.parameters(), lr=1e-2)
loss_fn = nn.BCEWithLogitsLoss()

n_epochs = 100
loss_history = []

customized_model.train()

for _ in range(n_epochs):
    optimizer.zero_grad()
    logits = customized_model(x_train)
    loss = loss_fn(logits, y_train_tensor)
    loss.backward()
    optimizer.step()
    loss_history.append(loss.item())

print("Loss before first update:", round(loss_history[0], 4))
print("Loss after final update: ", round(loss_history[-1], 4))

plt.figure(figsize=(6, 3.5))
plt.plot(range(1, n_epochs + 1), loss_history)
plt.xlabel("Epoch")
plt.ylabel("Binary cross-entropy loss")
plt.title("Training loss")
plt.tight_layout()
plt.show()

torch.manual_seed(0) x_train = e2t.align_features_to_input_nodes( data=x_train_df, artifact=artifact, ) x_test = e2t.align_features_to_input_nodes( data=x_test_df, artifact=artifact, ) customized_model = e2t.customize_model( model=model, head=nn.Linear(1, 1), ) y_train_tensor = torch.tensor( y_train.values.reshape(-1, 1), dtype=torch.float32, ) optimizer = torch.optim.Adam(customized_model.parameters(), lr=1e-2) loss_fn = nn.BCEWithLogitsLoss() n_epochs = 100 loss_history = [] customized_model.train() for _ in range(n_epochs): optimizer.zero_grad() logits = customized_model(x_train) loss = loss_fn(logits, y_train_tensor) loss.backward() optimizer.step() loss_history.append(loss.item()) print("Loss before first update:", round(loss_history[0], 4)) print("Loss after final update: ", round(loss_history[-1], 4)) plt.figure(figsize=(6, 3.5)) plt.plot(range(1, n_epochs + 1), loss_history) plt.xlabel("Epoch") plt.ylabel("Binary cross-entropy loss") plt.title("Training loss") plt.tight_layout() plt.show()

Loss before first update: 0.7286
Loss after final update:  0.1988

Interpret nodes¶

In this first walkthrough only input_signal is informative.

interpret_model() returns an InterpretationResult. For nodes, result.summary() is one table with hidden nodes only by default. Those scores are signed Captum values.

For state-update models, repeated node columns across steps are aggregated with aggregation="peak" by default (the signed score from the step with largest magnitude). Use result.sites to inspect per-step tables keyed by step_1, step_2, and so on. Plots below that use .abs() do that only for ranking.

In [9]:

node_result = e2t.interpret_model(
    model=customized_model,
    artifact=artifact,
    data=x_test_df,
    on="nodes",
    method="LayerConductance",
)
node_importance = node_result.summary()

node_importance.head()

node_result = e2t.interpret_model( model=customized_model, artifact=artifact, data=x_test_df, on="nodes", method="LayerConductance", ) node_importance = node_result.summary() node_importance.head()

Out[9]:

	node_a	node_b
0	-4.024604	0.0
1	-2.344724	0.0
2	-3.157850	0.0
3	-3.670888	0.0
4	-3.505065	0.0

In [10]:

node_attr_by_site = node_result.sites

list(node_attr_by_site.keys())

node_attr_by_site = node_result.sites list(node_attr_by_site.keys())

Out[10]:

['step_1', 'step_2']

In [11]:

for site_id, site_df in node_attr_by_site.items():
    print(site_id)
    display(site_df.head())

for site_id, site_df in node_attr_by_site.items(): print(site_id) display(site_df.head())

step_1

	input_noise	input_signal	node_a	node_b	output
0	0.0	0.0	-4.024604	0.0	0.0
1	0.0	0.0	-2.344724	0.0	0.0
2	0.0	0.0	-3.157850	0.0	0.0
3	0.0	0.0	-3.670888	0.0	0.0
4	0.0	0.0	-3.505065	0.0	0.0

step_2

	input_noise	input_signal	node_a	node_b	output
0	0.0	0.0	0.0	0.0	-4.024604
1	0.0	0.0	0.0	0.0	-2.344724
2	0.0	0.0	0.0	0.0	-3.157850
3	0.0	0.0	0.0	0.0	-3.670888
4	0.0	0.0	0.0	0.0	-3.505065

In [12]:

summary_a_only = (
    node_importance.abs()
    .mean(axis=0)
    .sort_values(ascending=True)
    .rename("mean_abs_attribution")
    .reset_index()
    .rename(columns={"index": "node"})
)

plt.figure(figsize=(7, 3.5))
sns.barplot(
    data=summary_a_only,
    x="mean_abs_attribution",
    y="node",
    color="steelblue",
)
plt.xlabel("Mean absolute node attribution")
plt.ylabel("Node")
plt.title("Summary node importance on the test set")
plt.tight_layout()
plt.show()

summary_a_only = ( node_importance.abs() .mean(axis=0) .sort_values(ascending=True) .rename("mean_abs_attribution") .reset_index() .rename(columns={"index": "node"}) ) plt.figure(figsize=(7, 3.5)) sns.barplot( data=summary_a_only, x="mean_abs_attribution", y="node", color="steelblue", ) plt.xlabel("Mean absolute node attribution") plt.ylabel("Node") plt.title("Summary node importance on the test set") plt.tight_layout() plt.show()

Positive and negative controls¶

The walkthrough above recovers importance on node_a when only input_signal is informative, while node_b stays near zero. That contrast is useful, but incomplete: it does not show what happens when the other input carries signal, when both do, or when neither does.

Hidden nodes and the output stay fixed. Each case rebuilds the edgelist and names inputs by role:

Case	Input → node_a	Input → node_b	Expected
A only	input_signal	input_noise	a high, b ~0
Both	input_signal_a	input_signal_b	both high
B only	input_noise	input_signal	both high
Neither	input_noise_a	input_noise_b	both ~0

A only reuses the walkthrough results above. The remaining cases use steps=3 so the feedback path input → node_b → node_a → output can reach the output. Code is omitted; each case shows the graph, training loss, and node importance on a shared attribution axis.

Case: A only¶

Only the input to node_a (input_signal) is informative. The input to node_b is input_noise. Results are taken from the walkthrough above.

Case: both inputs informative¶

Both inputs carry class signal (input_signal_a → node_a, input_signal_b → node_b). Both hidden nodes should receive importance.

Case: B only¶

Only the input to node_b (input_signal) is informative. The input to node_a is input_noise. Because only node_a connects to output, the path input_signal → node_b → node_a → output must carry the signal. Both hidden nodes should therefore receive importance.

Case: neither input informative¶

Both inputs are uninformative (input_noise_a, input_noise_b). The model should not recover strong, consistent importance on either hidden node.

Control summary¶

Across the four cases:

• A only (input_signal / input_noise): importance concentrates on node_a; node_b stays near zero
• Both (input_signal_a / input_signal_b): both hidden nodes receive importance
• B only (input_noise → node_a, input_signal → node_b): loss decreases and both hidden nodes receive importance via the feedback path to the output
• Neither (input_noise_a / input_noise_b): loss remains near chance; neither hidden node shows strong recovered importance

Low importance on node_b in the A-only setting therefore reflects its uninformative input. Recovering importance through the feedback path also requires enough unrolled steps for input → node_b → node_a → output to reach the output.