Recurrent example¶

This notebook walks through an end-to-end workflow for a graph with a feedback loop: train a small model and interpret named nodes.

It assumes you have already read the Getting started notebook. That tutorial covers package basics and a feedforward MaskedLinear network. Here we focus on what changes when the graph is not a DAG.

kpnn2 has two parsers, and you choose the layout:

• parse_layered() ranks a DAG into per-hop masks and skip records. It raises on a cycle, so it cannot parse this graph.
• parse_adjacency() puts every node into one state vector with packed source/target indices. Cycles and self-loops are allowed. That is the layout used here.

Recurrence itself stays ordinary PyTorch: one shared MaskedLinear over all nodes, applied for a fixed number of steps, with input values re-injected after each step. Weights are shared across steps. kpnn2 does not unroll time or pick the step count for you. This is not an LSTM, not a sequence model, and not a converge-until loop.

In [1]:

import random

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

import kpnn2

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

In [2]:

from matplotlib_inline.backend_inline import (
    set_matplotlib_formats,
)

random.seed(42)
np.random.seed(42)
torch.manual_seed(42)

set_matplotlib_formats(
    "svg",
    bbox_inches=None,
)
_FONT_SIZE = 14
_FIG_W = 7.0
_plot_rc = {
    "font.size": _FONT_SIZE,
    "axes.titlesize": _FONT_SIZE,
    "axes.labelsize": _FONT_SIZE,
    "xtick.labelsize": _FONT_SIZE,
    "ytick.labelsize": _FONT_SIZE,
    "legend.fontsize": _FONT_SIZE,
    "legend.title_fontsize": _FONT_SIZE,
    "lines.linewidth": 1.0,
    "font.family": "Liberation Sans",
    "text.color": "black",
    "axes.labelcolor": "black",
    "xtick.color": "black",
    "ytick.color": "black",
    "axes.edgecolor": "black",
    "axes.grid": False,
    "axes.facecolor": "none",
    "figure.facecolor": "none",
    "savefig.facecolor": "none",
    "savefig.edgecolor": "none",
    "svg.fonttype": "none",
}
plt.rcParams.update(_plot_rc)
sns.set_theme(
    style="ticks",
    rc=_plot_rc,
)


def _new_fig(height):
    return plt.figure(figsize=(_FIG_W, height))


def _style_ax(ax):
    ax.tick_params(labelsize=_FONT_SIZE)
    ax.title.set_fontsize(_FONT_SIZE)
    ax.xaxis.label.set_size(_FONT_SIZE)
    ax.yaxis.label.set_size(_FONT_SIZE)
    ax.set_facecolor("none")
    ax.grid(False)

from matplotlib_inline.backend_inline import ( set_matplotlib_formats, ) random.seed(42) np.random.seed(42) torch.manual_seed(42) set_matplotlib_formats( "svg", bbox_inches=None, ) _FONT_SIZE = 14 _FIG_W = 7.0 _plot_rc = { "font.size": _FONT_SIZE, "axes.titlesize": _FONT_SIZE, "axes.labelsize": _FONT_SIZE, "xtick.labelsize": _FONT_SIZE, "ytick.labelsize": _FONT_SIZE, "legend.fontsize": _FONT_SIZE, "legend.title_fontsize": _FONT_SIZE, "lines.linewidth": 1.0, "font.family": "Liberation Sans", "text.color": "black", "axes.labelcolor": "black", "xtick.color": "black", "ytick.color": "black", "axes.edgecolor": "black", "axes.grid": False, "axes.facecolor": "none", "figure.facecolor": "none", "savefig.facecolor": "none", "savefig.edgecolor": "none", "svg.fonttype": "none", } plt.rcParams.update(_plot_rc) sns.set_theme( style="ticks", rc=_plot_rc, ) def _new_fig(height): return plt.figure(figsize=(_FIG_W, height)) def _style_ax(ax): ax.tick_params(labelsize=_FONT_SIZE) ax.title.set_fontsize(_FONT_SIZE) ax.xaxis.label.set_size(_FONT_SIZE) ax.yaxis.label.set_size(_FONT_SIZE) ax.set_facecolor("none") ax.grid(False)

Define the graph¶

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 edgelist format is the one parse_layered() also takes: a source column and a target column, nothing else. parse_adjacency() additionally allows self-loops (a row where source == target), which this graph does not use.

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 [3]:

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[3]:

	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 [4]:

spec = kpnn2.parse_adjacency(edgelist)

spec.nodes, spec.input_nodes, spec.hidden_nodes, spec.output_nodes

spec = kpnn2.parse_adjacency(edgelist) spec.nodes, spec.input_nodes, spec.hidden_nodes, spec.output_nodes

Out[4]:

(('input_noise', 'input_signal', 'node_a', 'node_b', 'output'),
 ('input_noise', 'input_signal'),
 ('node_a', 'node_b'),
 ('output',))

In [5]:

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

for node_name in spec.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") for node_name in spec.nodes: dot.node(node_name, node_name) for row in edgelist.itertuples(index=False): dot.edge(str(row.source), str(row.target)) dot

Out[5]:

Shared masked update, unrolled for a fixed number of steps¶

parse_adjacency() already did the structural work:

1. named node sets, with inputs at in-degree 0 and outputs at out-degree 0: spec.input_nodes, spec.hidden_nodes, spec.output_nodes
2. packed source_index / target_index over spec.nodes; spec.to_mask() is the dense square for this small graph
3. where the inputs and outputs sit inside that state vector: spec.input_index, spec.output_index

The rest is ordinary PyTorch:

4. wrap spec.to_mask() in MaskedLinear
5. apply it for a fixed number of steps, re-injecting the input columns after each step
6. read the output nodes and apply a linear head

This notebook is small, so MaskedLinear is appropriate. See PackedLinear when n_nodes is large enough that the square would hurt RAM.

Re-injecting the inputs is required, not cosmetic. Input nodes have in-degree 0, so their rows of to_mask() are all zeros and their fan_in is 0. Under the degree-aware initialization of MaskedLinear those rows stay at zero forever, so nothing would ever write an input value into the state vector unless we do it ourselves.

Node interpretation sites are the unrolled steps (step_1, step_2). Each step is its own nn.Module so Captum can hook it, but they all share the same MaskedLinear weights. The first walkthrough names those two steps explicitly. Later controls add a step_3 so the feedback path can reach the output.

In [6]:

class RecurrentStep(nn.Module):
    def __init__(self, linear, input_index):
        super().__init__()
        # Shared MaskedLinear: do not register it as a child.
        object.__setattr__(self, "_linear", linear)
        self.input_index = input_index

    def forward(self, state, x):
        state = self._linear(state)
        state = state.clone()
        state[:, self.input_index] = x
        return state


class RecurrentNet(nn.Module):
    def __init__(self, spec: kpnn2.AdjacencySpec):
        super().__init__()
        self.spec = spec
        self.core = kpnn2.MaskedLinear(spec.to_mask())
        self.step_1 = RecurrentStep(
            self.core,
            spec.input_index,
        )
        self.step_2 = RecurrentStep(
            self.core,
            spec.input_index,
        )
        self.head = nn.Linear(1, 1)

    def forward(self, x):
        n_nodes = self.core.in_features
        state = x.new_zeros(x.shape[0], n_nodes)
        state[:, self.spec.input_index] = x
        state = self.step_1(state, x)
        state = self.step_2(state, x)
        output = state[:, self.spec.output_index]
        return self.head(output)

class RecurrentStep(nn.Module): def __init__(self, linear, input_index): super().__init__() # Shared MaskedLinear: do not register it as a child. object.__setattr__(self, "_linear", linear) self.input_index = input_index def forward(self, state, x): state = self._linear(state) state = state.clone() state[:, self.input_index] = x return state class RecurrentNet(nn.Module): def __init__(self, spec: kpnn2.AdjacencySpec): super().__init__() self.spec = spec self.core = kpnn2.MaskedLinear(spec.to_mask()) self.step_1 = RecurrentStep( self.core, spec.input_index, ) self.step_2 = RecurrentStep( self.core, spec.input_index, ) self.head = nn.Linear(1, 1) def forward(self, x): n_nodes = self.core.in_features state = x.new_zeros(x.shape[0], n_nodes) state[:, self.spec.input_index] = x state = self.step_1(state, x) state = self.step_2(state, x) output = state[:, self.spec.output_index] return self.head(output)

In [7]:

model = RecurrentNet(spec)
spec.nodes, spec.input_index, (model.step_1, model.step_2)

model = RecurrentNet(spec) spec.nodes, spec.input_index, (model.step_1, model.step_2)

Out[7]:

(('input_noise', 'input_signal', 'node_a', 'node_b', 'output'),
 (0, 1),
 (RecurrentStep(), RecurrentStep()))

Simulate data, align features, and train¶

In [8]:

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(name) for 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(name) for 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 [9]:

rng = np.random.default_rng(42)

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

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

x_train_df.head()

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

Out[9]:

	input_noise	input_signal
0	0.304717	-2.539984
1	0.750451	-0.559435
2	-1.951035	-2.802180
3	0.127840	-1.816243
4	-0.016801	-2.353044

In [10]:

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=list(spec.input_nodes),
    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)

n_cols = 2
g = sns.catplot(
    data=long_df,
    x="class",
    y="value",
    col="feature",
    col_wrap=n_cols,
    hue="class",
    kind="box",
    width=0.35,
    sharey=False,
    height=_FIG_W / n_cols,
    aspect=1.0,
    legend=False,
)
g.fig.set_size_inches(_FIG_W, g.fig.get_figheight())
g.fig.suptitle(
    "Simulated training data: informative vs. uninformative features",
    fontsize=_FONT_SIZE,
    y=0.98,
)
g.fig.subplots_adjust(
    top=0.82,
    bottom=0.16,
    left=0.1,
    right=0.98,
    wspace=0.3,
)
for ax, feature in zip(g.axes.flat, spec.input_nodes):
    feature_type = feature_type_map[feature]
    ax.set_title(f"{feature}\n({feature_type})")
    ax.set_xlabel("")
    ax.set_ylabel("value")
    _style_ax(ax)
g.fig.patch.set_facecolor("none")
plt.show()

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=list(spec.input_nodes), 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) n_cols = 2 g = sns.catplot( data=long_df, x="class", y="value", col="feature", col_wrap=n_cols, hue="class", kind="box", width=0.35, sharey=False, height=_FIG_W / n_cols, aspect=1.0, legend=False, ) g.fig.set_size_inches(_FIG_W, g.fig.get_figheight()) g.fig.suptitle( "Simulated training data: informative vs. uninformative features", fontsize=_FONT_SIZE, y=0.98, ) g.fig.subplots_adjust( top=0.82, bottom=0.16, left=0.1, right=0.98, wspace=0.3, ) for ax, feature in zip(g.axes.flat, spec.input_nodes): feature_type = feature_type_map[feature] ax.set_title(f"{feature}\n({feature_type})") ax.set_xlabel("") ax.set_ylabel("value") _style_ax(ax) g.fig.patch.set_facecolor("none") plt.show()

In [11]:

x_train = kpnn2.align_inputs(x_train_df, spec)
x_test = kpnn2.align_inputs(x_test_df, spec)

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

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=0.05,
)
loss_fn = nn.BCEWithLogitsLoss()

n_epochs = 100
loss_history = []
model.train()
for _ in range(n_epochs):
    optimizer.zero_grad()
    logits = 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),
)

_new_fig(2.8)
plt.plot(
    range(1, n_epochs + 1),
    loss_history,
    linewidth=1,
)
plt.xlabel("Epoch")
plt.ylabel("Binary cross-entropy loss")
plt.title("Training loss")
plt.tight_layout()
plt.show()

x_train = kpnn2.align_inputs(x_train_df, spec) x_test = kpnn2.align_inputs(x_test_df, spec) y_train_tensor = torch.tensor( y_train.values.reshape(-1, 1), dtype=torch.float32, ) optimizer = torch.optim.Adam( model.parameters(), lr=0.05, ) loss_fn = nn.BCEWithLogitsLoss() n_epochs = 100 loss_history = [] model.train() for _ in range(n_epochs): optimizer.zero_grad() logits = 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), ) _new_fig(2.8) plt.plot( range(1, n_epochs + 1), loss_history, linewidth=1, ) plt.xlabel("Epoch") plt.ylabel("Binary cross-entropy loss") plt.title("Training loss") plt.tight_layout() plt.show()

Loss before first update: 0.8023
Loss after final update:  0.1257

Interpret nodes¶

In this first walkthrough only input_signal is informative.

Captum LayerConductance is run on each unrolled step. Then map_node_attributions() labels the (batch, n_nodes) tensor as xarray with the full state-vector names. Convert with .to_pandas() for the tables below. kpnn2 does not import Captum. For the mapper itself, see Mapping attributions.

There is no layer= argument here. An AdjacencySpec has no depths, so the node axis is the whole state vector (spec.nodes) and the result carries no scalar layer coordinate.

Those scores are signed. For a summary like the old peak aggregation over steps, we keep the signed value from the step with largest magnitude, then take .abs() only for ranking. The default summary below is hidden nodes only (node_a, node_b). Use the per-step tables to inspect inputs and the output as well.

In [12]:

from captum.attr import LayerConductance

model.eval()
node_attr_by_site = {}
for site_id, step in (
    ("step_1", model.step_1),
    ("step_2", model.step_2),
):
    conductor = LayerConductance(model, step)
    attributions = conductor.attribute(
        x_test,
        target=0,
    )
    node_attr_by_site[site_id] = kpnn2.map_node_attributions(
        attributions=attributions,
        spec=spec,
    ).to_pandas()

list(node_attr_by_site.keys())

from captum.attr import LayerConductance model.eval() node_attr_by_site = {} for site_id, step in ( ("step_1", model.step_1), ("step_2", model.step_2), ): conductor = LayerConductance(model, step) attributions = conductor.attribute( x_test, target=0, ) node_attr_by_site[site_id] = kpnn2.map_node_attributions( attributions=attributions, spec=spec, ).to_pandas() list(node_attr_by_site.keys())

Out[12]:

['step_1', 'step_2']

In [13]:

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

node	input_noise	input_signal	node_a	node_b	output
observation
0	0.0	0.0	-5.821504	0.0	0.0
1	0.0	0.0	-2.579521	0.0	0.0
2	0.0	0.0	-2.648712	0.0	0.0
3	0.0	0.0	-5.395569	0.0	0.0
4	0.0	0.0	-6.949645	0.0	0.0

step_2

node	input_noise	input_signal	node_a	node_b	output
observation
0	0.0	0.0	0.0	0.0	-5.821504
1	0.0	0.0	0.0	0.0	-2.579521
2	0.0	0.0	0.0	0.0	-2.648712
3	0.0	0.0	0.0	0.0	-5.395569
4	0.0	0.0	0.0	0.0	-6.949644

In [14]:

def peak_hidden_summary(site_tables, hidden_nodes):
    stacked = np.stack(
        [table[hidden_nodes].to_numpy() for table in site_tables],
        axis=0,
    )
    pick = np.abs(stacked).argmax(axis=0)
    n_batch, n_hidden = pick.shape
    values = np.empty((n_batch, n_hidden))
    rows = np.arange(n_batch)
    for j in range(n_hidden):
        values[:, j] = stacked[
            pick[:, j],
            rows,
            j,
        ]
    return pd.DataFrame(
        values,
        columns=hidden_nodes,
    )


node_importance = peak_hidden_summary(
    list(node_attr_by_site.values()),
    list(spec.hidden_nodes),
)
node_importance.head()

def peak_hidden_summary(site_tables, hidden_nodes): stacked = np.stack( [table[hidden_nodes].to_numpy() for table in site_tables], axis=0, ) pick = np.abs(stacked).argmax(axis=0) n_batch, n_hidden = pick.shape values = np.empty((n_batch, n_hidden)) rows = np.arange(n_batch) for j in range(n_hidden): values[:, j] = stacked[ pick[:, j], rows, j, ] return pd.DataFrame( values, columns=hidden_nodes, ) node_importance = peak_hidden_summary( list(node_attr_by_site.values()), list(spec.hidden_nodes), ) node_importance.head()

Out[14]:

	node_a	node_b
0	-5.821504	0.0
1	-2.579521	0.0
2	-2.648712	0.0
3	-5.395569	0.0
4	-6.949645	0.0

In [15]:

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

_new_fig(2.4)
sns.barplot(
    data=summary_a_only,
    x="mean_abs_attribution",
    y="node",
    color="steelblue",
    width=0.3,
)
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"}) ) _new_fig(2.4) sns.barplot( data=summary_a_only, x="mean_abs_attribution", y="node", color="steelblue", width=0.3, ) 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 three named steps (step_1, step_2, step_3) so the feedback path input → node_b → node_a → output can reach the output.

In [16]:

def plot_graph(edgelist, input_a, input_b):
    dot = Digraph()
    dot.attr(rankdir="LR")
    for node_name in [input_a, input_b]:
        dot.node(node_name, node_name)
    for node_name in ["node_a", "node_b"]:
        dot.node(node_name, node_name)
    dot.node("output", "output")
    for row in edgelist.itertuples(index=False):
        dot.edge(
            str(row.source),
            str(row.target),
        )
    return dot


def make_edgelist(input_a, input_b):
    return pd.DataFrame(
        {
            "source": [
                input_a,
                input_b,
                "node_a",
                "node_b",
                "node_a",
            ],
            "target": [
                "node_a",
                "node_b",
                "node_b",
                "node_a",
                "output",
            ],
        }
    )


class RecurrentNet3(nn.Module):
    def __init__(self, spec: kpnn2.AdjacencySpec):
        super().__init__()
        self.spec = spec
        self.core = kpnn2.MaskedLinear(spec.to_mask())
        self.step_1 = RecurrentStep(
            self.core,
            spec.input_index,
        )
        self.step_2 = RecurrentStep(
            self.core,
            spec.input_index,
        )
        self.step_3 = RecurrentStep(
            self.core,
            spec.input_index,
        )
        self.head = nn.Linear(1, 1)

    def forward(self, x):
        n_nodes = self.core.in_features
        state = x.new_zeros(x.shape[0], n_nodes)
        state[:, self.spec.input_index] = x
        state = self.step_1(state, x)
        state = self.step_2(state, x)
        state = self.step_3(state, x)
        output = state[:, self.spec.output_index]
        return self.head(output)


def run_recurrent_case(
    edgelist,
    informative_features,
    n_epochs=100,
):
    torch.manual_seed(42)
    rng = np.random.default_rng(42)
    case_spec = kpnn2.parse_adjacency(edgelist)
    case_model = RecurrentNet3(case_spec)
    x_train_df, y_train = simulate_binary_data(
        n_per_class=80,
        feature_names=list(case_spec.input_nodes),
        informative_features=informative_features,
        rng=rng,
    )
    x_test_df, y_test = simulate_binary_data(
        n_per_class=40,
        feature_names=list(case_spec.input_nodes),
        informative_features=informative_features,
        rng=rng,
    )
    x_train = kpnn2.align_inputs(
        x_train_df,
        case_spec,
    )
    x_test = kpnn2.align_inputs(
        x_test_df,
        case_spec,
    )
    y_train_tensor = torch.tensor(
        y_train.values.reshape(-1, 1),
        dtype=torch.float32,
    )
    optimizer = torch.optim.Adam(
        case_model.parameters(),
        lr=0.05,
    )
    loss_fn = nn.BCEWithLogitsLoss()
    loss_history = []
    case_model.train()
    for _ in range(n_epochs):
        optimizer.zero_grad()
        logits = case_model(x_train)
        loss = loss_fn(logits, y_train_tensor)
        loss.backward()
        optimizer.step()
        loss_history.append(loss.item())
    case_model.eval()
    site_tables = []
    for step in (
        case_model.step_1,
        case_model.step_2,
        case_model.step_3,
    ):
        conductor = LayerConductance(
            case_model,
            step,
        )
        attributions = conductor.attribute(
            x_test,
            target=0,
        )
        site_tables.append(
            kpnn2.map_node_attributions(
                attributions=attributions,
                spec=case_spec,
            ).to_pandas()
        )
    peak = peak_hidden_summary(
        site_tables,
        list(case_spec.hidden_nodes),
    )
    summary = (
        peak.abs()
        .mean(axis=0)
        .sort_values(ascending=True)
        .rename("mean_abs_attribution")
        .reset_index()
        .rename(columns={"index": "node"})
    )
    return {
        "loss_history": loss_history,
        "summary": summary,
    }


def plot_case_results(
    case_name,
    result,
    attr_xlim=None,
):
    fig, axes = plt.subplots(
        1,
        2,
        figsize=(_FIG_W, 3.2),
        gridspec_kw={"width_ratios": [1.1, 1.0]},
    )
    axes[0].plot(
        range(1, len(result["loss_history"]) + 1),
        result["loss_history"],
        linewidth=1,
    )
    axes[0].set_xlabel("Epoch")
    axes[0].set_ylabel("Binary cross-entropy loss")
    axes[0].set_title(f"{case_name}: training loss")
    sns.barplot(
        data=result["summary"],
        x="mean_abs_attribution",
        y="node",
        color="steelblue",
        width=0.3,
        ax=axes[1],
    )
    axes[1].set_xlabel("Mean absolute node attribution")
    axes[1].set_ylabel("Node")
    axes[1].set_title(f"{case_name}: node importance")
    if attr_xlim is not None:
        axes[1].set_xlim(attr_xlim)
    _style_ax(axes[0])
    _style_ax(axes[1])
    fig.tight_layout()
    plt.show()


graph_a_only = plot_graph(
    edgelist,
    "input_signal",
    "input_noise",
)
result_a_only = {
    "loss_history": loss_history,
    "summary": summary_a_only,
}

edgelist_both = make_edgelist(
    "input_signal_a",
    "input_signal_b",
)
edgelist_b_only = make_edgelist(
    "input_noise",
    "input_signal",
)
edgelist_neither = make_edgelist(
    "input_noise_a",
    "input_noise_b",
)

graph_both = plot_graph(
    edgelist_both,
    "input_signal_a",
    "input_signal_b",
)
graph_b_only = plot_graph(
    edgelist_b_only,
    "input_noise",
    "input_signal",
)
graph_neither = plot_graph(
    edgelist_neither,
    "input_noise_a",
    "input_noise_b",
)

result_both = run_recurrent_case(
    edgelist=edgelist_both,
    informative_features=[
        "input_signal_a",
        "input_signal_b",
    ],
)
result_b_only = run_recurrent_case(
    edgelist=edgelist_b_only,
    informative_features=["input_signal"],
)
result_neither = run_recurrent_case(
    edgelist=edgelist_neither,
    informative_features=[],
)

xmax = max(
    result_a_only["summary"]["mean_abs_attribution"].max(),
    result_both["summary"]["mean_abs_attribution"].max(),
    result_b_only["summary"]["mean_abs_attribution"].max(),
    result_neither["summary"]["mean_abs_attribution"].max(),
)
attr_xlim = (0, float(xmax) * 1.05)

def plot_graph(edgelist, input_a, input_b): dot = Digraph() dot.attr(rankdir="LR") for node_name in [input_a, input_b]: dot.node(node_name, node_name) for node_name in ["node_a", "node_b"]: dot.node(node_name, node_name) dot.node("output", "output") for row in edgelist.itertuples(index=False): dot.edge( str(row.source), str(row.target), ) return dot def make_edgelist(input_a, input_b): return pd.DataFrame( { "source": [ input_a, input_b, "node_a", "node_b", "node_a", ], "target": [ "node_a", "node_b", "node_b", "node_a", "output", ], } ) class RecurrentNet3(nn.Module): def __init__(self, spec: kpnn2.AdjacencySpec): super().__init__() self.spec = spec self.core = kpnn2.MaskedLinear(spec.to_mask()) self.step_1 = RecurrentStep( self.core, spec.input_index, ) self.step_2 = RecurrentStep( self.core, spec.input_index, ) self.step_3 = RecurrentStep( self.core, spec.input_index, ) self.head = nn.Linear(1, 1) def forward(self, x): n_nodes = self.core.in_features state = x.new_zeros(x.shape[0], n_nodes) state[:, self.spec.input_index] = x state = self.step_1(state, x) state = self.step_2(state, x) state = self.step_3(state, x) output = state[:, self.spec.output_index] return self.head(output) def run_recurrent_case( edgelist, informative_features, n_epochs=100, ): torch.manual_seed(42) rng = np.random.default_rng(42) case_spec = kpnn2.parse_adjacency(edgelist) case_model = RecurrentNet3(case_spec) x_train_df, y_train = simulate_binary_data( n_per_class=80, feature_names=list(case_spec.input_nodes), informative_features=informative_features, rng=rng, ) x_test_df, y_test = simulate_binary_data( n_per_class=40, feature_names=list(case_spec.input_nodes), informative_features=informative_features, rng=rng, ) x_train = kpnn2.align_inputs( x_train_df, case_spec, ) x_test = kpnn2.align_inputs( x_test_df, case_spec, ) y_train_tensor = torch.tensor( y_train.values.reshape(-1, 1), dtype=torch.float32, ) optimizer = torch.optim.Adam( case_model.parameters(), lr=0.05, ) loss_fn = nn.BCEWithLogitsLoss() loss_history = [] case_model.train() for _ in range(n_epochs): optimizer.zero_grad() logits = case_model(x_train) loss = loss_fn(logits, y_train_tensor) loss.backward() optimizer.step() loss_history.append(loss.item()) case_model.eval() site_tables = [] for step in ( case_model.step_1, case_model.step_2, case_model.step_3, ): conductor = LayerConductance( case_model, step, ) attributions = conductor.attribute( x_test, target=0, ) site_tables.append( kpnn2.map_node_attributions( attributions=attributions, spec=case_spec, ).to_pandas() ) peak = peak_hidden_summary( site_tables, list(case_spec.hidden_nodes), ) summary = ( peak.abs() .mean(axis=0) .sort_values(ascending=True) .rename("mean_abs_attribution") .reset_index() .rename(columns={"index": "node"}) ) return { "loss_history": loss_history, "summary": summary, } def plot_case_results( case_name, result, attr_xlim=None, ): fig, axes = plt.subplots( 1, 2, figsize=(_FIG_W, 3.2), gridspec_kw={"width_ratios": [1.1, 1.0]}, ) axes[0].plot( range(1, len(result["loss_history"]) + 1), result["loss_history"], linewidth=1, ) axes[0].set_xlabel("Epoch") axes[0].set_ylabel("Binary cross-entropy loss") axes[0].set_title(f"{case_name}: training loss") sns.barplot( data=result["summary"], x="mean_abs_attribution", y="node", color="steelblue", width=0.3, ax=axes[1], ) axes[1].set_xlabel("Mean absolute node attribution") axes[1].set_ylabel("Node") axes[1].set_title(f"{case_name}: node importance") if attr_xlim is not None: axes[1].set_xlim(attr_xlim) _style_ax(axes[0]) _style_ax(axes[1]) fig.tight_layout() plt.show() graph_a_only = plot_graph( edgelist, "input_signal", "input_noise", ) result_a_only = { "loss_history": loss_history, "summary": summary_a_only, } edgelist_both = make_edgelist( "input_signal_a", "input_signal_b", ) edgelist_b_only = make_edgelist( "input_noise", "input_signal", ) edgelist_neither = make_edgelist( "input_noise_a", "input_noise_b", ) graph_both = plot_graph( edgelist_both, "input_signal_a", "input_signal_b", ) graph_b_only = plot_graph( edgelist_b_only, "input_noise", "input_signal", ) graph_neither = plot_graph( edgelist_neither, "input_noise_a", "input_noise_b", ) result_both = run_recurrent_case( edgelist=edgelist_both, informative_features=[ "input_signal_a", "input_signal_b", ], ) result_b_only = run_recurrent_case( edgelist=edgelist_b_only, informative_features=["input_signal"], ) result_neither = run_recurrent_case( edgelist=edgelist_neither, informative_features=[], ) xmax = max( result_a_only["summary"]["mean_abs_attribution"].max(), result_both["summary"]["mean_abs_attribution"].max(), result_b_only["summary"]["mean_abs_attribution"].max(), result_neither["summary"]["mean_abs_attribution"].max(), ) attr_xlim = (0, float(xmax) * 1.05)

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.

In [17]:

display(graph_a_only)
plot_case_results("A only", result_a_only, attr_xlim)

display(graph_a_only) plot_case_results("A only", result_a_only, attr_xlim)

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.

In [18]:

display(graph_both)
plot_case_results(
    "Both informative",
    result_both,
    attr_xlim,
)

display(graph_both) plot_case_results( "Both informative", result_both, attr_xlim, )

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.

In [19]:

display(graph_b_only)
plot_case_results("B only", result_b_only, attr_xlim)

display(graph_b_only) plot_case_results("B only", result_b_only, attr_xlim)

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.

In [20]:

display(graph_neither)
plot_case_results(
    "Neither informative",
    result_neither,
    attr_xlim,
)

display(graph_neither) plot_case_results( "Neither informative", result_neither, attr_xlim, )

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.