How to compute the penalty for invariant risk minimization in Tensorflow?

user_10575373 0
Viewed 284

I am trying to implement the technique called "Invariant risk minimization," which adds a penalty term to the loss function in training machine learning models. The new penalty term's technical definition is the squared gradient norm with respect to a constant classifier. There is an implementation of this "penalty" function with PyTorch here.

I was wondering how I can implement this function in Tensorflow 2.

More specifically, I want to implement the function below, which is also in the code I shared its link.

  def penalty(logits, y):
    scale = torch.tensor(1.).cuda().requires_grad_()
    loss = mean_nll(logits * scale, y)
    grad = autograd.grad(loss, [scale], create_graph=True)[0]
    return torch.sum(grad**2)
loss-function tensorflow machine-learning python
1 Answers

In a similar fashion :

def penalty(y_true, y_pred):
    scale = tf.constant(1.)
    with tf.GradientTape() as tape:
        tape.watch(scale)
        loss = tf.losses.binary_crossentropy(y_true, y_pred*scale, from_logits=True)
    grad = tape.gradient(loss, [scale])[0]
    return tf.reduce_sum(grad**2)

Note that the order of the parameter logits and ground truth are reversed compared to the PyTorch version, to respect TensorFlow's convention.

To calculate a Gradient, you just need tf.GradientTape, you can read more in the guide : Introduction to Gradients and Automatic Differentiation

Comparing that the 2 versions produces the same results :

PyTorch version:

>>> penalty(torch.tensor(1.),torch.tensor(0.))
tensor(0.5344, grad_fn=<SumBackward0>)

>>> penalty(torch.tensor(1.),torch.tensor(1.))
tensor(0.0723, grad_fn=<SumBackward0>)

TensorFlow version:

>>> penalty([0.],[1.])
<tf.Tensor: shape=(), dtype=float32, numpy=0.53444666>

>>> penalty([1.],[1.])
<tf.Tensor: shape=(), dtype=float32, numpy=0.0723295>
user_73701530
Related