Unrolling backprop through update (maml & nested grad ops)

[choltz95]

Hi, I am trying to learn training sample weights by doing maml on a validation set:

So I want to do backprop through a one-step update - to compute the gradients of training sample weights w.r.t the validation loss:

def maml_loss(batch_weights, params, batch_a, batch_b):
    x1, y1 = batch_a  # training set
    x2, y2 = batch_b # validation set
    
    # one-step 'inner' update on network parameters using training set
    for _ in range(1):
        grads = grad(loss)(params, batch_weights, batch_a) # compute param grads of weighted training loss
        p = tree_multimap(lambda x,y: x-y, p, grads) # inner update

    # compute unweighted validation loss on updated params
    # want to compute training weight grads of this, p is a fn of the gradient of the one-step loss, 
    # the one-step loss is a fn of the training sample weights
    return loss(params, batch_b) 

# want to recover gradients of sample weights - just returns zeros for some reason
grad(maml_loss)(batch_weights, ....) 

the weighted training loss looks like this:

def loss(...):
    logits = stax.logsoftmax(logits) 
    return -np.sum(np.sum(logits * targets, axis=1)*ex_weights)  # weighted cross entropy loss

The issue is that I'm recovering sample weight gradients that are just zero. I would guess that this should be doable, since nesting grad ops & doing 'normal' maml on regular parameter weights works just fine.

I guess the training weights are being treated as constants at some point?

[choltz95]

Embarrassed to say that it looks like I miss-broadcasted things.
Feel free to remove this.