Model does not learn any variance

[e-mauss]

Hello everyone,
i am currently facing a problem where my model is learning nothing besides the mean of my target values.
Im training this model for a regression problem where i will have a lot of graphs as samples in my dataset. These graphs will always have the same topology, as they are always modeling the same graph-like object, just at a different point in time. Given that circumstance, edge features and some node features will only differ across the different nodes/edges inside a graph, but not across the many graphs themselves (i am unsure if it even makes sense to keep those features then).

Further example on this im not sure how to express it without it being confusing

Some features, such as feature c will never differ across all the graphs but only across nodes. Features such as d may differ without such constraints.

graph number | node A          | node B 
             | feat c | feat d | feat c | feat d
_____________|________|________|________|________
           1 |   0.4  |   0.38 |   0.26 |   0.99 
           2 |   0.4  |   0.55 |   0.26 |   0.44
           3 |   0.4  |   0.56 |   0.26 |   0.38
           4 |   0.4  |   0.78 |   0.26 |   0.45

I am trying to predict two values (magnitude and phase) for every node in every graph/timestep of a target validation set.

From the following loss and accuracy plots it does seem like my model is learning. (A prediction is considered accurate when both magnitude and phase are within a given margin of error from the target value.

The plots

It even seems to predict the phase quite well, however i think thats just because theres not a lot of variance in the targets to begin with.

Phase violin plot

When it comes to voltage magnitudes there is very little to no variance at all in the predictions. That can also be seen in the following boxplot where the predictions should be shown as a yellow bar to the right of every green bar. However, due to the lack of variance there are no yellow bars (or whiskers) to be seen.

Magnitude violin plot and boxplot

When plotting my predictions and targets for one particular node of every graph you can see my model doesnt seem to learn the underlying relations. The predictions are plotted in green, the targets are plotted in orange.

Preds and targets

It seems that the ~80% accuracy achieved by my model are owed to around 80% of the targets being with the currently specified margins of error from the predicted values.

All of the above plots originate from the following training run:

The training run

Optimizer: Adam, learning rate: 0.178
Scheduler: CosineAnnealingLR, T_max: 35
Epochs: 35
n_layers (param for model below): 40
hidden_channels (param for model below): 64
batch_size: 32

According to this stackexchange post it seems like my model is not complex enough and gives up on learning any variance, instead just settling for always predicting the mean as the best way to minimize the output of the loss function. However I am unsure how to increase the complexity of a GNN.

My model

class PowerflowNet2(tnn.Module):
    def __init__(self, hidden_channels, out_channels, n_layers, norm):
        super(PowerflowNet2, self).__init__()
        self.node_encoder = Linear(-1, hidden_channels)
        self.edge_encoder = Linear(-1, hidden_channels)
        self.lin = Linear(hidden_channels, out_channels)
        self.n_layers = n_layers

        self.convs = tnn.ModuleList()
        self.acts = tnn.ModuleList()
        self.norms = tnn.ModuleList()
        for layer in range(n_layers):
            self.convs.append(GENConv(hidden_channels, hidden_channels, aggr='add', norm=norm))
            self.acts.append(tnn.ReLU(inplace=True))
            self.norms.append(tnn.LayerNorm(hidden_channels, elementwise_affine=True))

    def forward(self, x, edge_index, edge_attr):

        n_layers = self.n_layers

        x = self.node_encoder(x)
        edge_attr = self.edge_encoder(edge_attr)

        for layer in range(n_layers):
            x = self.convs[layer](x, edge_index, edge_attr)
            x = self.norms[layer](x)
            x = self.acts[layer](x)

        x = self.lin(x)

        return x

I have also found /discussions/4740 which seems to have dealt with the exact same problem, however there seems to be no real conclusion on what to do to improve the model.

As I have little to no experience with graph neural networks i thought i might just ask for some help here.
Thank you very much in advance,
Erik :)

[rusty1s]

Wow, what a detailed question :) I am not totally sure I can help you but if you are dealing with the problem that all your model is learning is the mean of the targets, you should try out some weighted loss formulation, i.e., weight node targets higher that are far away from the mean.