LibAUC: A Deep Learning Library for X-Risk Optimization
Zhuoning Yuan Dixian Zhu Zi-Hao Qiu Gang Li Xuanhui Wang Tianbao Yang
This paper introduces the award-winning deep learning (DL) library called
LibAUC for implementing state-of-the-art algorithms towards optimizing a family
of risk functions named X-risks. X-risks refer to a family of compositional
functions in which the loss function of each data point is defined in a way
that contrasts the data point with a large number of others. They have broad
applications in AI for solving classical and emerging problems, including but
not limited to classification for imbalanced data (CID), learning to rank
(LTR), and contrastive learning of representations (CLR). The motivation of
developing LibAUC is to address the convergence issues of existing libraries
for solving these problems. In particular, existing libraries may not converge
or require very large mini-batch sizes in order to attain good performance for
these problems, due to the usage of the standard mini-batch technique in the
empirical risk minimization (ERM) framework. Our library is for deep X-risk
optimization (DXO) that has achieved great success in solving a variety of
tasks for CID, LTR and CLR. The contributions of this paper include: (1) It
introduces a new mini-batch based pipeline for implementing DXO algorithms,
which differs from existing DL pipeline in the design of controlled data
samplers and dynamic mini-batch losses; (2) It provides extensive benchmarking
experiments for ablation studies and comparison with existing libraries. The
LibAUC library features scalable performance for millions of items to be
contrasted, faster and better convergence than existing libraries for
optimizing X-risks, seamless PyTorch deployment and versatile APIs for various
loss optimization. Our library is available to the open source community at
https://github.com/Optimization-AI/LibAUC, to facilitate further academic
research and industrial applications.