A relationship type identifies what kind of connection an edge represents. In a molecular graph, edges can represent single, double, triple, or aromatic bonds. A graph with more than one edge type is heterogeneous; a graph with only one is homogeneous. The distinction matters when a model must handle different connections differently rather than treating every edge alike. section 5.4.3
Relationship types enter graph models in two practical ways:
- When types are known, pass an edge-type value alongside node features and edge indices. A heterogeneous graph convolution uses parameters specific to each relation, then aggregates messages by summation. Several such layers can form an encoder for molecular graph generation. section 5.4.3
- When connections are unknown, infer their types. In neural relational inference, an encoder produces probabilities over edge types; sampled types guide a decoder that predicts future node trajectories. The example uses two types—connected and not connected—but allows more. section 6.4 section 6.5.3
Listing 5.19 Heterogenous GCN
from torch.nn import Parameter
from torch_geometric.nn import MessagePassing
class HeterogeneousGraphConv(MessagePassing):
def __init__(self, in_channels, out_channels, num_relations, bias=True):
super(HeterogeneousGraphConv, self).\
__init__(aggr='add') #1
self.in_channels = in_channels
self.out_channels = out_channels
self.num_relations = num_relations
self.weight = Parameter(torch.\
Tensor(num_relations, in_channels, \
out_channels)) #2
if bias:
self.bias = Parameter(torch.Tensor(out_channels))
else:
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
torch.nn.init.xavier_uniform_(self.weight)
if self.bias is not None:
torch.nn.init.zeros_(self.bias)
def forward(self, x, edge_index, edge_type):
return self.propagate\
(edge_index, size=(x.size(0),
x.size(0)), x=x, edge_type=edge_type) #3
def message(self, x_j, edge_type, index, size): #4
W = self.weight[edge_type] #5
x_j = torch.matmul(x_j.unsqueeze(1), W).squeeze(1)
return x_j
def update(self, aggr_out):
if self.bias is not None:
aggr_out += self.bias
return aggr_out
The PyTorch Geometric message-passing layer allocates a weight tensor with a separate slice for each relation and selects summation for aggregation. section 5.4.3
Listing 6.11 Loss for the NRI model
def calculate_loss(self, inputs,
is_train=False,
teacher_forcing=True,
return_edges=False,
return_logits=False):
encoder_results = self.encoder(inputs)
logits = encoder_results['logits']
hard_sample = not is_train
edges = F.gumbel_softmax\
(logits.view(-1, self.num_edge_types),
tau=self.gumbel_temp,
hard=hard_sample).view\
(logits.shape) #1
output = self.decoder(inputs[:, :-1], edges)
if len(inputs.shape) == 3: \
target = inputs[:, 1:]
else:
Target = inputs[:, 1:, :, :]
loss_nll = F.mse_loss(\
output, target) / (2 * \
self.prior_variance) #2
probs = F.softmax(logits, dim=-1)
log_probs = torch.log(probs + 1e-16) #3
loss_kl = (probs * \
(log_probs - torch.log(\
torch.tensor(1.0 / #4
self.num_edge_types)))).\
sum(-1).mean()
loss = loss_nll + loss_kl
return loss, loss_nll, loss_kl, logits, output
The neural relational inference loss shows Gumbel-Softmax sampling from edge-type logits before the decoder uses the resulting edges. section 6.4
Choose the treatment to match the data. Known bond categories can be supplied to a heterogeneous encoder. If the graph structure itself is uncertain, edge types become predictions: a prior expresses the starting assumption about connectivity, and Gumbel temperature affects discrete sampling. For temporal predictions, the decoder also needs to capture how the graph changes over time. section 5.4.3 section 6.4 section 6.5.3
