Kpnn2: sparsely connected PyTorch layers from a named graph, with attributions by node name

projects
1

kpnn2 turns a graph of named entities into sparsely connected PyTorch layers. Each node becomes a unit, each edge becomes a weight, pairs without an edge stay unconnected, and attribution scores come back labeled by node name.

kpnn2_overview
kpnn2_overview3307×1163 236 KB

You provide an edgelist and data with named features, and you train in ordinary PyTorch.

From edgelist to layers

import pandas as pd
import kpnn2

edgelist = pd.DataFrame(
    [("input_signal_1", "hidden_signal"), ("input_signal_2", "hidden_signal"),
     ("input_noise_1", "hidden_noise"), ("input_noise_2", "hidden_noise"),
     ("hidden_signal", "output"), ("hidden_noise", "output")],
    columns=["source", "target"],
)
spec = kpnn2.parse_layered(edgelist)
for layer in spec.layer_nodes:
    print(layer)

('input_noise_1', 'input_noise_2', 'input_signal_1', 'input_signal_2')
('hidden_noise', 'hidden_signal')
('output',)

The edgelist describes this network, and the three printed layers are its columns:

network
network1446×688 31.3 KB

Plain PyTorch modules

PackedLinear stores one trainable weight per edge. identity=spec.fingerprint makes a checkpoint from a different graph refuse to load.

import torch
from torch import nn

torch.manual_seed(42)
hop_0, hop_1 = spec.hops  # a hop: all edges entering one layer
model = nn.Sequential(
    kpnn2.PackedLinear.from_hop(hop_0, identity=spec.fingerprint),
    nn.Tanh(),
    kpnn2.PackedLinear.from_hop(hop_1, identity=spec.fingerprint),
)

Data matched by name

The labels depend only on the two signal inputs.

names = ["input_signal_1", "input_signal_2", "input_noise_1", "input_noise_2"]
features = pd.DataFrame(torch.randn(200, 4).numpy(), columns=names)
labels = features["input_signal_1"] + features["input_signal_2"] > 0

col = kpnn2.align_inputs(features.columns, spec)  # match columns by name
x = torch.as_tensor(features.to_numpy()[:, col], dtype=torch.float32)
y = torch.as_tensor(labels.to_numpy(), dtype=torch.float32).unsqueeze(1)

optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
for _ in range(200):
    optimizer.zero_grad()
    nn.functional.binary_cross_entropy_with_logits(model(x), y).backward()
    optimizer.step()

Attributions by node name

Any attribution method works. This one uses Captum (pip install captum), which kpnn2 does not depend on.

from captum.attr import LayerConductance

attributions = LayerConductance(model, model[0]).attribute(x)
scores = kpnn2.map_node_attributions(attributions, spec, hop_output=hop_0)
print(abs(scores).mean("observation").to_pandas().round(2))

node
hidden_noise     0.39
hidden_signal    4.90
dtype: float32

Why not plain PyTorch?

You can mask dense weights, but then you write the parser yourself, and inside the model a node name is only a tensor position. When the two drift apart (reordered columns, an updated graph, a reloaded checkpoint), training still runs, but scores land under the wrong names. kpnn2 checks the match wherever names meet tensor positions and raises an error instead. The model stays your own PyTorch code, and cyclic, recurrent, and attention models work too.

Why not a GNN (graph NN) with PyG?

In a GNN the graph is the data. In kpnn2 the graph is the architecture: one fixed graph shared by all samples, with its own weights on every edge. More details are on the kpnn2 website (see link below).

Where does this apply?

Wherever the nodes are known entities and the edges are known relationships between them, for example in biology or chemistry. Such models are called knowledge-primed neural networks (KPNNs), the “kpnn” in kpnn2. The point is interpretability: every hidden unit belongs to a known entity, such as a signaling protein in a cell or a chemical species in a reaction network, so attribution scores point to things you can test in the lab, not to anonymous units.

Links

Feedback is very welcome!