Including global graph information into the loss function

[ognjenkundacina]

Dear ignnition team,

In our problem we are training a GNN for a regression task - a subset of nodes is labeled by a float value and those labels are learned using a neural network as a readout model. The whole GNN model is trained based on the MSE loss between the labels and the predictions for the mentioned subset of nodes and works well!

We would like to create a new loss function that incorporates some physical laws related to our problem. In each training step, after the predictions are generated for all of the labeled nodes, we would like to add an additional term to the MSE between the labels and the predictions. That additional term would multiply all of the predictions with some coefficients (different for every node), and sum all of the obtained values. So the goal would be to minimize that sum along with the MSE.

Is something like this possible to implement in the ignnition framework? I'm not even sure is it consistent with the logic for creating the mini-batches - I guess that the requirement here would be to have all of the nodes from one training sample in the same mini-batch.

Thanks!

Ognjen

[ognjenkundacina]

Hi José,

Thank you for the quick response! I'll be slowly working on this in the next two months, so I suppose I will not have any questions soon. If it's convenient for you, you can close the issue.

Regards,

Ognjen

[jsuarezv]

Hi Ognjen,
What you want to do is supported by IGNNITION, although it requires some non-standard workarounds. Here you can find an example that does something similar:
https://github.com/BNN-UPC/ignnition/tree/main/examples/Radio_resource_allocation

In general, you can follow these steps:

1. Add the loss coefficients as node-level features in the dataset (e.g., $wmmse_power in the example).
2. Define a “concat” function at the end of the readout that includes the final predictions and the coefficients. This can be done in the model_description.yaml file. Note that in the wireless network example the readout function is not implemented with a neural network. If you want to add a neural network there you may use the following code snippet:

readout:
- type: neural_network
  input: [transmitter_receiver_pair]
  nn_name: readout
  output_name: readout_value
- type: concat
  axis: 1
  input: [readout_value, $wmmse_power]
  output_label: [$path_loss]

Here, "readout_value" would be the actual output of the readout function, while the concatenated [readout_value, $wmmse_power] list is just used for the loss function.

3. Define a separate file for additional functions and a custom name for the loss function in the train_options.yaml file:

additional_functions_file: ./additional_functions.py
...
loss: sum_rate_loss

4. In the additional_functions.py file you may define a custom loss function (e.g., “sum_rate_loss” in the example). This function must have “y_true” and “y_pred” as input arguments, as in the example. In this function, “y_true” are the training labels defined as “output_label” in the model_description.yaml file (i.e., $path_loss in the example), while “y_pred” is a list that contains the tensors previously defined in the "concat" function. In the example we define [readout_value, $wmmse_power] as "input" to the concat function (model_description.yaml; line 37). Then, you can access the two features as follows:

wmmse_power = tf.expand_dims(y_pred[:, -1], axis=0)
readout_value = tf.expand_dims(y_pred[:, -2], axis=0)

Here, "wmmse_power" would be the equivalent of the coefficients you would use in your problem, and "readout_value" the predictions made by the model.

Note that a similar approach can be followed to set a custom evaluation metric. For example here we define the custom “sum_rate_metric”. Then, in the “additional_functions_file.py” we implement this metric.

Also, note that during inference the coefficients must also be defined in the dataset, although they will not be used by the model, so this feature can be simply set to arbitrary values.

I hope you find this explanation helpful. Please let us know if you have further questions.

Regards,
José