How to correctly update parameters wrt to other model

[averageFlaxUser]

Hello, I am trying to validate my implementation of a bi-level optimization implementation I have. My goal is the following:

"Some model_A that produces output vector z, which is used as a regularization term in the loss_fn of some model_B that is trained on some data. I am looking to optimize the params of model_A with respect to the performance of model_B."

My pseudo-Jax/Flax implementation:

def aug_cross_entropy_loss(*, logits, labels, z):
    labels_onehot = jax.nn.one_hot(labels, num_classes=10)
    softmax_logits = softmax(logits)
    nll = jnp.log(softmax_logits)
    nll = jnp.multiply(nll, labels_onehot)
    nll = jnp.multiply(nll, z) # some regularization being done
    loss = -jnp.mean(nll)
    return loss

def eval_fn(model_B, input_to_B):
    logits = model_B.apply_fn({'params': model_B.params}, input_to_B['image'])
    loss = optax.softmax_cross_entropy_with_integer_labels(
        logits=logits, labels=input_to_B['label']).mean()
    return loss

def inner_step(model_B, input_to_B, z):
    def loss_fn(params_B):
        logits = model_B.apply_fn({'params': params_B}, input_to_B['image'])
        loss = aug_cross_entropy_loss(logits=logits, labels=input_to_B['label'], z=z)
        return loss
    grad_fn = jax.value_and_grad(loss_fn)
    for i in range(5): # perform K-gradient steps
        loss, grads = grad_fn(model_B.params)
        model_B = model_B.apply_gradients(grads=grads)
    return model_B

def outer_step(model_A, model_B, input_to_A, input_to_B):
    def outer_loss(params_A):
        z = model_A.apply_fn({'params': params_A}, input_to_A)
        new_model_B = inner_step(model_B, input_to_B, z)
        return eval_fn(new_ model_B, batch), new_model_B
    grad_fn = jax.value_and_grad(outer_loss, has_aux=True)
    (loss, model_B), grads = grad_fn(model_A.params)
    model_A = model_A.apply_gradients(grads=grads)
    return model_A, model_B

Train loop, model initialization and other items are done the same way as in the Flax Getting Started Guide.

[cgarciae]

At least in this example I think it should be

grad_fn = jax.value_and_grad(outer_loss, has_aux=True)
(loss, model_B), grads = grad_fn(model_A.params)
...

You need to reassign model_B else changes are not being propagated.

[averageFlaxUser]

Yes, thats a typo on my behalf. Will fix it in the original post.

[cgarciae]

Makes sense! It would help if you can post a minimal example we can run.

[averageFlaxUser]

Yes, let me get that! Thanks.

[averageFlaxUser]

@cgarciae Here is the notebook. Couple notes, the example works but it might not make "sense" what the overall goal is and thats fine. I just want to figure out if I am optimizing the meta_state correctly w.r.t. the performance of the state model.