2 losses, 1 forward pass

[charleswilmot]

Hello everyone,

I can't find documentation for my use case, maybe because I'm not sure how to formulate my problem.

Say I have a networks composed of 10 layers, with weights W_0, W_1, ..., W_9.

Now I want to optimize these weights with respect to 2 loss functions called loss_outer and loss_inner, where

• loss_inner only applies to weights W_4, W_5, W_6
• loss_outer applies to weights W_0, ..., W_3 and W_7, ..., W_9

@jax.jit
def loss_inner(inner_weights, outer_weights, x_4):
    x_5 = layer_4(W_4, x_4)
    x_6 = layer_5(W_5, x_5)
    x_7 = layer_6(W_6, x_6)
    return something(x_4, ..., x_7)

@jax.jit
def loss_outer(outer_weights, inner_weights, x_0):
    x_1 = layer_0(W_0, x_0)
    x_2 = layer_1(W_1, x_1)
    x_3 = layer_2(W_2, x_2)
    x_4 = layer_3(W_3, x_3)
    # so called 'inner' part
    x_5 = layer_4(W_4, x_4)
    x_6 = layer_5(W_5, x_5)
    x_7 = layer_6(W_6, x_6)
    #
    x_8 = layer_7(W_7, x_7)
    x_9 = layer_8(W_8, x_8)
    x_out = layer_9(W_9, x_9)
    return something(x_0, ..., x_out)

@jax.jit
def gradients(network_weights, x_0):
    # separate the weights in 2 partitions, according to the text above
    inner_weights, outer_weights = partition(network_weights)
    x_1 = layer_0(W_0, x_0)
    x_2 = layer_1(W_1, x_1)
    x_3 = layer_2(W_2, x_2)
    x_4 = layer_3(W_3, x_3)
    
    inner_grad = jax.grad(inner_loss)(inner_weights, outer_weights, x_4)
    outer_grad = jax.grad(outer_loss)(outer_weights, inner_weights, x_0)

    # does the inverse of the 'partition' function
    gradients = recombine(inner_grad, outer_grad)
    return gradients

In the code above, will XLA take care of simplifying the computation graph such that the forward pass is computed only once?

Can you think of a better way to compute the gradients, without computing the forward pass twice ?

Many thanks to anyone who's willing to help me with that :-)

[YouJiacheng]

You can check the optimized IR.
If you really want to stop the gradient of inner loss to outer weight(I don't know why), maybe you can try jax.lax.stop_gradient.
Or try partial eval, see #9778.
Implement a custom vjp might be suitable as well.