Batch representation orthnormalization, gini regulariztion

[rhsimplex]

Hi @rusty1s,

I'm interested in adding two regularizers to pytg that we proposed in our recently accepted ICML paper: Improving Molecular Graph Neural Network Explainability with Orthonormalization and Induced Sparsity.

These are both used like L1 or L2 weight decay, though typically only for specific layers. So they're not normalization layers, probably best just implemented as functions.

Before I make a PR, just wanted to 1) make sure it would be likely to be accepted and 2) discuss where something like this might go...I was thinking either torch_geometric.nn.norm or torch_geometric.utils

Thanks!

[rusty1s]

Thanks for the pointer :) Yes, I'm definitely interested in this PR!

I think such regularization should have its place in torch_geometric.nn, and we can probably start thinking about adding a sub-package torch_geometric.nn.functional to the library. WDYT?

How does the API of such regularization techniques will look like?

[rhsimplex]

Yep, that makes sense to me. I'll prepare a PR then and we can discuss.

The signatures look something like:

def gini_score(w):
     """w is a weight matrix"""
     ...
     return score  # float [0, 1]

def bro_penalty(x, batch):
    """x, batch are the usual pytg x and batch"""
    ...
    return penalty # float [0, inf)