minimum_irm.py

"""This is directly from https://arxiv.org/pdf/1907.02893.pdf"""
import torch
from torch.autograd import grad


def compute_penalty(losses, dummy):
  g1 = grad(losses[0::2].mean(), dummy, create_graph=True)[0]
  g2 = grad(losses[1::2].mean(), dummy, create_graph=True)[0]
  return (g1 * g2).sum()


def example_1(n=10000, d=2, env=1):
  x = torch.randn(n, d) * env
  y = x + torch.randn(n, d) * env
  z = y + torch.randn(n, d)
  return torch.cat((x, z), 1), y.sum(1, keepdim=True)


phi = torch.nn.Parameter(torch.ones(4, 1))
dummy_w = torch.nn.Parameter(torch.Tensor([1.0]))
opt = torch.optim.SGD([phi], lr=1e-3)
mse = torch.nn.MSELoss(reduction="none")
environments = [example_1(env=0.1), example_1(env=1.0)]
for iteration in range(50000):
  error = 0
  penalty = 0
  for x_e , y_e in environments:
    p = torch.randperm(len(x_e))
    error_e = mse(x_e[p] @ phi * dummy_w, y_e [p])
    penalty += compute_penalty(error_e, dummy_w)
    error += error_e.mean()
  opt.zero_grad()
  (1e-5 * error + penalty).backward()
  opt.step()
  if iteration % 1000 == 0:
    print(phi)
    print('erm error', 1e-5 * error.data.cpu().numpy())
    print('grad penalty', penalty.data.cpu().numpy())