Train a neural network when the training has only the derivative of output wrt all inputs

Question

There is a scalar function F with 1000 inputs. I want to train a model to predict F given the inputs. However, in the training dataset, we only know the derivative of F with respect to each input, not the value of F itself. How I can construct a neural network with this limitation in tensorflow or pytorch?

Answer 1

I think you can use torch.autograd to compute the gradients, and then use them for the loss. You need:

(a) A trainable nn.Module to represent the (unknown) function F:

class UnknownF(nn.Module):
  def __init__(self, ...):
    # whatever combinations of linear layers and activations and whatever...

  def forward(self, x):
    # x is 1000 dim vector
    y = self.layers(x)
    # y is a _scalar_ output
    return y

model = UnknownF(...)  # instansiate the model of the unknown function

(b) Training data:

x = torch.randn(n, 1000, requires_grad=True)  # n examples of 1000-dim vectors
dy = torch.randn(n, 1000)  # the corresponding n-dim gradients of the n inputs

(c) An optimizer:

opt = torch.optim.SGD(model.parameters(), lr=0.1)

(d) Put it together:

criterion = nn.MSELoss()

for e in range(num_epochs):
  for i in range(n):
    # batch size = 1, pick one example
    x_ = x[i, :]
    dy_ = dy[i, :] 
    opt.zero_grad()
    # predict the unknown output
    y_ = model(x_)
    # compute the gradients of the model using autograd:
    pred_dy_ = autograd.grad(y_, x_, create_graph=True)[0]
    # compute the loss between the model's gradients and the GT ones:
    loss = criterion(pred_dy_, dy_)
    loss.backward()
    opt.step()  # update model's parameters accordingly.