arXiv:2410.21231cs.LGcs.MS2024-10JMLR被引 8

用Wasserstein距离提升模型鲁棒性,代码改动极小即可实现

$\texttt{skwdro}$: a library for Wasserstein distributionally robust machine learning

  • 基于Wasserstein距离的分布鲁棒优化,通过熵平滑增强灵活性
  • 仅需少量代码修改即可为PyTorch模型添加鲁棒训练能力
  • 兼容scikit-learn接口,适合希望快速部署鲁棒学习的开发者

我们提出skwdro,一个用于训练鲁棒机器学习模型的Python库。该库基于基于Wasserstein距离的分布鲁棒优化,广泛应用于最优传输与机器学习领域。其目标是通过为PyTorch模块提供封装器,使更多用户能以最小代码改动实现模型损失的鲁棒化。库中还包含若干常见目标的scikit-learn兼容估计器。核心实现采用原始鲁棒目标的熵平滑方法,以确保最大模型灵活性。代码与文档分别托管于https://github.com/iutzeler/skwdro和https://skwdro.readthedocs.io。

原文摘要 · Abstract (English)

We present skwdro, a Python library for training robust machine learning models. The library is based on distributionally robust optimization using Wasserstein distances, popular in optimal transport and machine learnings. The goal of the library is to make the training of robust models easier for a wide audience by proposing a wrapper for PyTorch modules, enabling model loss' robustification with minimal code changes. It comes along with scikit-learn compatible estimators for some popular objectives. The core of the implementation relies on an entropic smoothing of the original robust objective, in order to ensure maximal model flexibility. The library is available at https://github.com/iutzeler/skwdro and the documentation at https://skwdro.readthedocs.io.

鲁棒学习WassersteinPyTorch

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。