Skip to content
Go to file

Latest commit


Git stats


Failed to load latest commit information.
Latest commit message
Commit time
Jun 22, 2020
Oct 6, 2020
Jun 22, 2020
Oct 5, 2020
Jun 25, 2020
Oct 6, 2020
Jun 22, 2020
Jun 22, 2020

AuxiLearn - Auxiliary Learning by Implicit Differentiation

This repository contains the source code to support the paper Auxiliary Learning by Implicit Differentiation, by Aviv Navon*, Idan Achituve*, Haggai Maron, Gal Chechik and Ethan Fetaya.


  1. Paper
  2. Project page


Clone the repo and install:

git clone
cd AuxiLearn
pip install .


Given a bi-level optimization problem in which the upper-level parameters (i.e., auxiliary parameters) are only implicitly affecting the upper-level objective, you can use auxilearn to compute the upper-level gradients through implicit differentiation.

The main code component you will need to use is auxilearn.optim.MetaOptimizer. It is a wrapper over PyTorch optimizers that updates its parameters through implicit differentiation.

Code example

We assume two models, primary_model and auxiliary_model, and two dataloaders. The primary_model is optimized using the train data in the train_loader, and the auxiliary_model is optimized using the auxiliary set in the aux_loader. We assume a loss_fuction that return the train loss if train=True, or auxiliary set loss if train=False. Also, we assume the training loss is a function of both the primary parameters and the auxiliary parameters, and that the loss on the auxiliary set (or validation set) is a function of the primary parameters only. In Auxiliary Learning, the auxiliary set loss is the loss on the main task (see paper for more details).

from auxilearn.optim import MetaOptimizer

primary_model = MyModel()
auxiliary_model = MyAuxiliaryModel()
# optimizers
primary_optimizer = torch.optim.Adam(primary_model.parameters())

aux_lr = 1e-4
aux_base_optimizer = torch.optim.Adam(auxiliary_model.parameters(), lr=aux_lr)
aux_optimizer = MetaOptimizer(aux_base_optimizer, hpo_lr=aux_lr)

# training loop
step = 0
for epoch in range(epochs):
    for batch in train_loder:
        step += 1
        # calculate batch loss using 'primary_model' and 'auxiliary_model'
        loss = loss_func(train=True)
        # update primary parameters
        # condition for updating auxiliary parameters
        if step % aux_params_update_every == 0:
            # calc current train loss
            train_set_loss = loss_func(train=True)
            # calc current auxiliary set loss - this is the loss over the main task
            auxiliary_set_loss = loss_func(train=False) 
            # update auxiliary parameters - no need to call loss.backwards() or aux_optimizer.zero_grad()


If you find auxilearn to be useful in your own research, please consider citing the following paper:

  title={Auxiliary Learning by Implicit Differentiation},
  author={Navon, Aviv and Achituve, Idan and Maron, Haggai and Chechik, Gal and Fetaya, Ethan},
  journal={arXiv preprint arXiv:2007.02693},
You can’t perform that action at this time.