Scatter version of Cross-Entropy Loss

[leonardcaquot94]

Context

I’m working on a node selection task where I compute logits for each node in a batch and train them using cross-entropy loss and selection probabilities. Since I select one node per graph and the number of nodes per graph varies, cross-entropy must be computed per graph using the batch tensor for grouping.

I propose adding a scatter-based cross-entropy function, similar to scatter_softmax in torch_geometric.utils, to compute one loss per graph.

Implementation

import torch
from torch_geometric.utils import scatter, softmax

def scatter_cross_entropy(logits, target, batch, reduce='mean'):
    log_probs = torch.log(softmax(logits, batch))
    return -scatter(target * log_probs, batch, reduce=reduce)

Reduction

The reduce parameter mirrors the behavior of reduction parameter in torch.nn.functional.cross_entropy to compute a per-graph loss. An additional reduction can be applied outside this function to obtain a single scalar value if needed.