DOCSEARCHbyAlgolia
Submit Docs

LibAUC: Deep Learning Library for X-Risk Optimization

repository·main·Indexed 43 days ago

https://github.com/optimization-ai/libauc

LibAUC is a deep learning library specialized in X-Risk (extreme risk) optimization, providing APIs to directly optimize performance measures such as AUROC, AUPRC, and NDCG. It supports imbalanced classification (CID), learning to rank (LTR), and contrastive learning of representation (CLR). The library includes specialized losses like AUCMLoss and NDCGLoss, and optimizers such as PESG and SONG, with LibAUC 2.0 introducing a Trainer, distributed training support, and Hugging Face integration.

Tokens
10.3K
Snippets
34
Records
40
Agent score
55%

What's inside LibAUC

  1. Overview of LibAUC capabilities

    main

    LibAUC is a deep learning library designed for X-Risk Optimization. It provides a unified framework to optimize compositional loss functions across several domains:

    • Classification of Imbalanced Data (CID): Surrogate losses for AUROC, AUPRC/AP, and partial AUROC.
    • Learning to Rank (LTR): Surrogate losses for NDCG, top-K NDCG, and listwise losses.
    • Contrastive Learning of Representation (CLR): Global contrastive losses.

    LibAUC 2.0 introduces the LibAUC Trainer, distributed training support for all optimizers, and integration with Hugging Face models.

  2. Install LibAUC via pip or source

    main

    You can install LibAUC using pip for a quick setup, or clone the repository from GitHub to install from source.

    Using pip:

    $ pip install -U libauc

    From source:

    $ git clone https://github.com/Optimization-AI/LibAUC.git
$ cd LibAUC
$ pip install .
    pip install -U libauc

  3. How SogCLR with Cosine Gamma Schedule works

    main

    In bimodal contrastive learning tasks using the SogCLR algorithm, the parameter $\gamma$ (the coefficient for the moving average estimator) is typically set to a constant. This tutorial demonstrates using a cosine $\gamma$ decay schedule instead.

    Given:

    • $e$: current epoch
    • $E$: number of decay epochs (recommended to set to half of the total training epochs)
    • $\gamma_{\mathrm{min}}$: the final value

    The value of $\gamma$ at epoch $e$ is calculated as: $$\gamma = (1.0- \gamma_{\mathrm{min}}) * 0.5 * \cos(\pi * e / E) + \gamma_{\mathrm{min}}$$

    If $e > E$, $\gamma$ is set to $\gamma_{\mathrm{min}}$. This schedule helps SogCLR outperform standard CLIP models by improving early-stage training performance.

  4. How to optimize NDCG using SONG and K-SONG

    main

    LibAUC provides specialized optimizers and loss functions for optimizing NDCG (Normalized Discounted Cumulative Gain). The recommended workflow involves two stages:

    1. Initial Warm-up: Train the model using ListwiseCELoss and the SONG optimizer. This optimizes the listwise cross-entropy loss to provide a good starting point.
    2. Main Optimization: Switch to NDCG_Loss (either standard SONG or top-K K-SONG) to optimize the target NDCG metric directly.

    Key Components

    • Losses: ListwiseCELoss (for warm-up) and NDCG_Loss (for main optimization).
    • Optimizers: SONG is used for both stages.
    • Samplers: TriSampler is used to handle the ranking task by sampling positive and negative items.
    • Models: The example uses NeuMF (Neural Matrix Factorization), but these components are compatible with other architectures.
    # Warm-up stage
criterion = ListwiseCELoss(N=num_relevant_pairs, num_pos=NUM_POS, gamma=GAMMA0)
optimizer = SONG(model.parameters(), lr=LR, weight_decay=L2, mode=OPTIMIZER_STYLE)

# Main optimization stage (SONG or K-SONG)
criterion = NDCGLoss(num_relevant_pairs, n_users, n_items, NUM_POS, gamma0=GAMMA0, topk=TOPK, topk_version='theo')
optimizer = SONG(params=model.parameters(), lr=LR, weight_decay=L2, mode=OPTIMIZER_STYLE)

  5. Optimize NDCG using SONG and K-SONG

    main

    LibAUC provides specialized optimizers and loss functions for optimizing NDCG (Normalized Discounted Cumulative Gain) in ranking tasks.

    SONG (Stochastic Optimization of NDCG)

    SONG is an algorithm for NDCG optimization with provable convergence. It uses the NDCG_Loss class and the SONG optimizer.

    K-SONG

    K-SONG is an algorithm for top-K NDCG optimization. It is implemented using the same NDCG_Loss and SONG classes but requires specific parameters for top-K optimization.

    Key Hyperparameters

    • lr: Initial learning rate (highly sensitive).
    • gamma0: Moving average factor for the algorithm (highly sensitive). Recommended range: (0.0, 1.0). Default is 0.1.
    • beta1: Momentum parameter (similar to Adam, typically 0.9).
    • topk_version: For K-SONG, use 'theo' for theoretical K-SONG or 'prac' for practical K-SONG.
    # Example setup for K-SONG
from libauc.losses.ranking import NDCG_Loss
from libauc.optimizers import SONG

# Parameters
TOPK = 300
TOPK_V = 'theo' # 'prac' or 'theo'
GAMMA0 = 0.1

# Build optimizer and criterion
criterion = NDCG_Loss(id_mapper, num_relevant_pairs, n_users, n_items, NUM_POS, gamma0=GAMMA0, topk=TOPK, topk_version=TOPK_V)
optimizer = SONG(params=model.parameters(), lr=LR, weight_decay=L2, mode='adam')

  6. Optimize Listwise Cross-Entropy using Warm-up

    main

    For better performance in NDCG optimization, it is recommended to perform an initial warm-up phase by optimizing the Listwise Cross-Entropy loss before switching to SONG or K-SONG.

    Implementation

    • Loss: Use ListwiseCE_Loss.
    • Optimizer: Use SONG.

    Parameters

    • gamma0: Moving average factor (highly sensitive).
    from libauc.losses.ranking import ListwiseCE_Loss
from libauc.optimizers import SONG

# Build optimizer and criterion for warm-up
criterion = ListwiseCE_Loss(id_mapper=id_mapper, total_relevant_pairs=num_relevant_pairs, num_pos=NUM_POS, gamma0=GAMMA0)
optimizer = SONG(params=model.parameters(), lr=LR, weight_decay=L2, mode='adam')

  7. Perform Linear Evaluation on Pretrained Features

    main

    After pretraining a model using contrastive loss, you can evaluate its performance by performing linear classification on the frozen features.

    Steps:

    1. Load the pretrained checkpoint.
    2. Load the base encoder (e.g., resnet18) but replace the projection head with a linear layer (nn.Linear) for the target number of classes.
    3. Load the state_dict using strict=False to skip the non-linear projection layers that were only used during pretraining.
    4. Freeze all parameters except the new linear layer (param.requires_grad = False).
    5. Train only the linear layer using standard CrossEntropyLoss and an optimizer like torch.optim.SGD.
    # 1. Load model and checkpoint
model = resnet18(pretrained=False, last_activation=None, num_classes=128)
# ... (setup model architecture) ...

checkpoint = torch.load(checkpoint_dir, map_location="cpu")
state_dict = checkpoint['state_dict']

# 2. Filter out projection head weights (e.g., 'fc')
for k in list(state_dict.keys()):
    if 'fc' in k:
       del state_dict[k]
model.load_state_dict(state_dict, strict=False)

# 3. Freeze backbone
for name, param in model.named_parameters():
    if name not in ['fc.weight', 'fc.bias']:
       param.requires_grad = False

# 4. Train only the linear head
parameters = list(filter(lambda p: p.requires_grad, model.parameters()))
optimizer = torch.optim.SGD(parameters, lr=lr, momentum=0.9)

  8. Setup Dataset Pipeline for Contrastive Learning

    main

    Contrastive learning requires a specific data pipeline where each image is transformed into two different augmented crops (pairwise samples).

    When using libauc.datasets.CIFAR10, ensure you set return_index=True to retrieve the image index along with the image and label, as this is required by the GCLoss function to identify samples.

    Workflow:

    1. Define a TwoCropsTransform class that applies two different augmentations to the same input.
    2. Use libauc.datasets.CIFAR10(..., return_index=True, transform=...).
    3. The training loop will then receive (images, _, index) where images contains the two crops and index is the sample identifier.
    # Example of the expected data structure in the training loop
for i, (images, _, index) in enumerate(train_loader):
    # images[0] is the first crop, images[1] is the second crop
    # index is used for the loss function
    hidden1 = model(images[0].cuda())
    hidden2 = model(images[1].cuda())
    loss = loss_fn(hidden1, hidden2, index)

  10. Import LibAUC components for Contrastive Learning

    main

    To perform self-supervised training using Global Contrastive Loss (GCLoss) and SogCLR, import the following modules:

    import libauc
from libauc.models import resnet50, resnet18
from libauc.datasets import CIFAR100
from libauc.optimizers import SogCLR, LARS
from libauc.losses import GCLoss
    import libauc
from libauc.models import resnet50, resnet18
from libauc.datasets import CIFAR100
from libauc.optimizers import SogCLR, LARS
from libauc.losses import GCLoss

import torch
import torchvision.transforms as transforms
import torch.nn as nn
import numpy as np
import os,math,shutil

  11. Adjust Gamma during training

    main

    When using a dynamic gamma schedule, you must manually trigger the gamma adjustment at the start of each epoch within your training loop. This is done by calling the adjust_gamma method on the model's criterion object.

    # Inside the training loop, at the start of each epoch:
if hasattr(model_orig.criterion, 'adjust_gamma'):
    model_orig.criterion.adjust_gamma(epoch)

  12. Extreme Classification Tutorial: SCENT and SOX

    main

    This tutorial demonstrates how to perform extreme classification using Compositional Entropic Risk (CERM) via the SCENT and SOX algorithms. It uses a subset of the TreeOfLife-10M biology dataset (163K species) and implements a linear classifier in PyTorch.

    Workflow Overview:

    1. Data Preparation: Download pre-computed features.
    2. Model Definition: Define a LinearClassifier that maps feature vectors to logits for $m$ classes.
    3. Optimization: Use libauc.optimizers.SCENT and libauc.losses.EntLossClassification.
    4. Training Loop: Implement a standard PyTorch training loop, utilizing criterion.adjust_gamma(epoch, epochs) if available to update risk parameters over time.
    5. Evaluation: Monitor validation accuracy and cross-entropy loss to save the best model checkpoints.