callous-youth/BOAT

A Compositional Operation Toolbox for Gradient-based Bi-Level Optimization

What it solves

BOAT simplifies the implementation of Bi-Level Optimization (BLO), which is often complex to move from theoretical mathematical models to practical code. Instead of providing a few fixed solver routines, it provides a modular system that allows researchers to build and experiment with a vast variety of BLO solvers without rewriting their core model code.

How it works

BOAT factorizes the BLO workflow into atomic, reusable gradient operations. It organizes these into three modular libraries:

  1. Gradient Mapping (GM-OL): Reconstructs the iterative trajectory of the lower-level problem.
  2. Numerical Approximation (NA-OL): Handles the bottleneck of calculating implicit or hyper-gradients using methods like automatic differentiation or numerical inversion.
  3. First-Order (FO-OL): Reformulates nested problems into single-level objectives to avoid expensive Hessian computations.

Users define their optimization strategy and objectives via JSON configuration files. BOAT then automatically composes these primitives into a functional solver (supporting over 85 variants from 19 atomic operations) and integrates them with PyTorch models and optimizers.

Who it’s for

It is designed for AI researchers and practitioners working on tasks that require nested optimization, such as meta-learning, hyperparameter optimization, and neural architecture search.

Highlights

  • Compositional Design: Deconstructs solvers into modular stages for flexible assembly.
  • Generative Solvers: Allows the creation of novel hybrid solvers by simply changing a configuration file.
  • Configuration-Driven: Decouples algorithmic logic from model definitions using JSON.
  • Broad Application Support: Applicable to data-centric (e.g., data cleaning), model-centric (e.g.,PEFT), and strategy-centric (e.g., RLHF) AI tasks.
  • Multi-Backend Support: While primarily PyTorch-based, it also offers versions for Jittor and MindSpore.

Related

  • Project
    metaopt/torchopt

    An efficient PyTorch library for differentiable optimization that supports explicit, implicit, and zero-order differentiation for bi-level optimization and meta-learning.

  • Project
    mosaicml/composer

    A PyTorch-based deep learning library that simplifies distributed training at scale, providing tools for memory optimization and algorithmic speedups to reduce training time and cost.

  • Project
    kyegomez/zeta

    A modular PyTorch framework providing reusable, high-performance building blocks for assembling state-of-the-art AI models, including attention mechanisms, quantization, and multi-modal architectures.

  • Project
    pypose/bae

    A PyTorch-based library for exact 2nd-order sparse optimization in robotics, specifically designed to accelerate Bundle Adjustment and Pose Graph Optimization using CUDA.

  • Project
    pypose/pypose

    A PyTorch-based library that combines deep perceptual models with physics-based optimization for robot learning, providing tools for Lie groups, filtering, and second-order optimization.