[Feature Proposal] Topology-Constrained Linear Layer (Watts–Strogatz / Barabási–Albert weight masks)

[narwanRanak]

Summary

I'd like to contribute a new module — TopologyMaskedLinear — that constrains the active weight positions of a dense linear layer to follow a graph-theoretic topology (Watts–Strogatz small-world or Barabási–Albert scale-free). This fills a gap in PyG's nn.dense module: while PyG has powerful GNN convolution layers that operate on externally provided graphs, there is currently no primitive for constraining the internal weight connectivity of a linear layer using graph topology.

Contributors: Ranak (ranak23426@iiitd.ac.in) — Primary · Kartikeya (kartikeya22241@iiitd.ac.in) — Secondary

---

Motivation

Research in network science has shown that neural network weight matrices, when thresholded, naturally develop small-world and scale-free properties during training (e.g., Watts & Strogatz 1998, Barabási & Albert 1999). A natural question is: what happens if we impose these topologies as structural inductive biases at initialisation?

We have been running experiments on CIFAR-10/100 using WS and BA graph masks on MLP weight matrices, tracking how the topology evolves during training and whether the small-world structure is preserved. The main challenges are:

1. Gradient erosion — BA hub edges accumulate larger gradients, eroding WS ring-edge magnitudes over time.
2. No standard primitive — building these masks and keeping them attached to a nn.Module requires boilerplate every time.

A reusable PyG module would make this class of experiment much cleaner to express.

---

Proposed API

Graph mask generators (in torch_geometric.utils)

from torch_geometric.utils.graph_masks import (
    watts_strogatz_mask,               # → BoolTensor (num_rows, num_cols)
    watts_strogatz_mask_with_ring,     # → (BoolTensor, BoolTensor) — mask + original ring
    barabasi_albert_mask,              # → BoolTensor (num_rows, num_cols)
)

These are pure-PyTorch bipartite adaptations of the classical WS and BA models — no NetworkX dependency.

TopologyMaskedLinear (in torch_geometric.nn.dense)

A drop-in replacement for torch_geometric.nn.dense.Linear with a fixed topology mask:

from torch_geometric.nn.dense import TopologyMaskedLinear

# Watts–Strogatz small-world mask
layer = TopologyMaskedLinear(512, 512, graph_type='ws', k=6, p=0.3, seed=42)

# Barabási–Albert scale-free mask
layer = TopologyMaskedLinear(512, 512, graph_type='ba', m=3, seed=42)

# Or bring your own mask
from torch_geometric.utils.graph_masks import barabasi_albert_mask
mask = barabasi_albert_mask(512, 512, m=3)
layer = TopologyMaskedLinear(512, 512, mask=mask)

# Forward — masked positions zeroed out, only active edges learned
x = torch.randn(32, 512)
out = layer(x)  # shape: (32, 512)
print(layer.sparsity)  # e.g. 0.988

The mask is registered as a non-trainable buffer. Forward: F.linear(x, weight * mask, bias).

WSEdgeRegLoss (in torch_geometric.nn.dense)

An optional regularisation loss that prevents gradient descent from eroding the WS ring structure:

from torch_geometric.nn.dense import WSEdgeRegLoss
from torch_geometric.utils.graph_masks import watts_strogatz_mask_with_ring

mask, ring_mask = watts_strogatz_mask_with_ring(512, 512, k=6, p=0.3, seed=42)
layer = TopologyMaskedLinear(512, 512, mask=mask)
reg_loss = WSEdgeRegLoss(target=1.2)

# In training loop:
total_loss = criterion(output, labels) + 0.01 * reg_loss(layer, ring_mask)

---

Implementation plan

We have a working implementation ready. Proposed file locations:

torch_geometric/
  utils/
    graph_masks.py                       ← ws/ba mask generators (pure PyTorch, no NetworkX)
  nn/
    dense/
      topology_masked_linear.py          ← TopologyMaskedLinear + WSEdgeRegLoss
test/
  utils/
    test_graph_masks.py
  nn/
    dense/
      test_topology_masked_linear.py

Happy to open a PR once there's a green light on the API design.

---

Prior work / references

• Watts, D.J. & Strogatz, S.H. (1998). Collective dynamics of small-world networks. Nature, 393, 440–442.
• Barabási, A.L. & Albert, R. (1999). Emergence of scaling in random networks. Science, 286(5439), 509–512.
• Experiments on WS-BA hybrid graph networks for CIFAR-10/100 classification — IIIT Delhi, Network Science research (Ranak, primary; Kartikeya, secondary).

---

Questions for the team

1. Is torch_geometric.nn.dense the right home, or would you prefer torch_geometric.nn.sparse or a new submodule?
2. Should the mask generators live in torch_geometric.utils or closer to torch_geometric.nn?
3. Is there appetite for WSEdgeRegLoss, or should I keep that in a follow-up PR?

Happy to adjust scope based on feedback. Thanks for considering this!