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

Home | Installation | Docs | Tutorials | Examples

BOAT is a compositional OperAtion-level Toolbox for gradient-based BLO. Unlike existing libraries that typically encapsulate fixed solver routines, BOAT factorizes the BLO workflow into atomic, reusable primitives. Through a unified constraint reconstruction perspective, it empowers researchers to automatically compose over 85+ solver variants from a compact set of 17 gradient operations.

This is the Jittor-based version of BOAT (boat-jit). It leverages Jittor’s Just-In-Time (JIT) compilation and efficient CUDA/cuDNN backends to significantly accelerate large-scale gradient-based BLO experiments.

BOAT supports unified execution across backends via separate branches:

• PyTorch-based: The widely-used standard version.
• MindSpore-based: Optimized for ascending chips and industry scenarios.

🔑 Key Features

• 🧩 Compositional Operation-Level Abstraction: Deconstructs BLO solvers into three modular stages: Gradient Mapping (GM), Numerical Approximation (NA), and First-Order (FO).
• ⚡ Accelerated JIT Execution: Built on Jittor, enabling meta-operator fusion and high-performance execution on NVIDIA GPUs.
• 🏭 Generative Solver Construction: Supports dynamic serialization of operations. Users can recover classical algorithms or discover novel hybrid solvers simply by changing configurations.
• 🛠 Configuration-Driven: Define complex optimization strategies via simple JSON configurations, decoupling algorithmic logic from model definitions.
• ✅ Comprehensive Testing: Achieves 99% code coverage through rigorous testing with pytest, ensuring software robustness.

📚 Supported Operation Libraries

🔨 Installation

BOAT-jit is built on top of Jittor. Please ensure Jittor is installed correctly before installing BOAT.

1. Install Jittor

Follow the Official Installation Guide or use the commands below:

Linux (Ubuntu / CentOS)

sudo apt install python3.8-dev libomp-dev
python3.8 -m pip install jittor
# Verify installation (Optional)
python3.8 -m jittor.test.test_example

Windows

python -m pip install jittor
python -m jittor.test.test_core

macOS

brew install libomp
python3.8 -m pip install jittor

2. Install BOAT-jit

Once Jittor is ready, install BOAT-jit via PyPI or Source:

# Install from PyPI
pip install boat-jit

# Or install from Source (Specific Branch)
git clone -b boat_jit --single-branch https://github.com/callous-youth/BOAT.git
cd BOAT
pip install -e .

Note: examples/L2_Reg_jit additionally requires scikit-learn (pip install scikit-learn).

✅ Quick Verification

Run the following commands from the root directory of the BOAT repository:

python examples/L2_Reg_jit/l2_regularization.py

python examples/neural_architecture_search/neural_architecture_search.py

⚡ How to Use BOAT

BOAT separates the problem definition from the solver configuration. Below is a Jittor-based example.

1. Load Configurations

import json
import boat_jit as boat

# Load algorithmic configurations
with open("configs/boat_config.json", "r") as f:
    boat_config = json.load(f)

# Load objective configurations
with open("configs/loss_config.json", "r") as f:
    loss_config = json.load(f)

2. Define Models and Optimizers

Use standard Jittor models and optimizers.

import jittor as jit

# Define models
upper_model = MyUpperModel()
lower_model = MyLowerModel()

# Define optimizers (Jittor syntax)
upper_opt = jit.nn.Adam(upper_model.parameters(), lr=1e-3)
lower_opt = jit.nn.SGD(lower_model.parameters(), lr=1e-2)

3. Customize & Initialize Problem

Inject runtime objects into the configuration.

# Configure BOAT with Jittor models/optimizers
boat_config["lower_level_model"] = lower_model
boat_config["upper_level_model"] = upper_model
boat_config["lower_level_opt"] = lower_opt
boat_config["upper_level_opt"] = upper_opt

# Initialize the BOAT core
b_optimizer = boat.Problem(boat_config, loss_config)

4. Build Solvers

Inject runtime objects into the configuration.

# Pass optimizers explicitly if needed by the solver backend
b_optimizer.build_ll_solver(lower_opt)
b_optimizer.build_ul_solver(upper_opt)

5. Run Optimization Loop

# Training loop
for x_itr in range(1000):
    # Prepare data (Jittor Arrays or dicts)
    ul_feed_dict = {"data": ul_data, "target": ul_target}
    ll_feed_dict = {"data": ll_data, "target": ll_target}
    
    # Run one iteration
    loss, run_time = b_optimizer.run_iter(ll_feed_dict, ul_feed_dict, current_iter=x_itr)
    
    if x_itr % 100 == 0:
        print(f"Iter {x_itr}: UL Loss {loss:.4f}")

🌍 Applications

BOAT covers a wide spectrum of BLO applications, categorized by the optimization target:

• Data-Centric: Data Hyper-Cleaning, Synthetic Data Reweighting, Diffusion Model Guidance.
• Model-Centric: Neural Architecture Search (NAS), LLM Prompt Optimization, Parameter Efficient Fine-Tuning (PEFT).
• Strategy-Centric: Meta-Learning, Hyperparameter Optimization (HO), Reinforcement Learning from Human Feedback (RLHF).

📝 Citation

If you find BOAT useful in your research, please consider citing our paper:

@article{liu2025boat,
  title={BOAT: A Compositional Operation Toolbox for Gradient-based Bi-Level Optimization},
  author={Liu, Yaohua and Pan, Jibao and Jiao, Xianghao and Gao, Jiaxin and Liu, Zhu and Liu, Risheng},
  journal={Submitted to Journal of Machine Learning Research (JMLR)},
  year={2025}
}

License

MIT License

Copyright (c) 2024 Yaohua Liu

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.