JAXopt: 硬件加速、可批处理和可微分的JAX优化器

jaxopt

JAXopt简介

JAXopt是一个基于JAX的优化器库,为各种优化问题提供了高效、灵活的解决方案。它具有以下主要特点:

硬件加速: JAXopt的实现可以在GPU和TPU上运行,充分利用硬件加速能力。
可批处理: 使用JAX的vmap功能,可以自动向量化多个相同优化问题的实例。
可微分: 优化问题的解可以相对于输入进行隐式微分或通过展开算法迭代的自动微分。

这些特性使JAXopt成为一个强大的优化工具,适用于各种科学计算和机器学习任务。

JAXopt Logo

安装

安装JAXopt非常简单,可以通过pip直接安装最新发布版:

pip install jaxopt

如果需要安装开发版本,可以使用以下命令:

pip install git+https://github.com/google/jaxopt

主要功能

JAXopt提供了多种优化算法的实现,包括:

无约束优化:
- 梯度下降法
- BFGS算法
- L-BFGS算法
约束优化:
- 投影梯度下降法
- 增广拉格朗日法
二次规划:
- OSQP求解器
- BoxOSQP求解器
非光滑优化:
- 近似梯度法
- 邻近梯度法
随机优化:
- SGD
- Adam
- RMSProp
根查找:
- 牛顿法
- Halley法
非线性最小二乘:
- 高斯-牛顿法
- Levenberg-Marquardt算法

这些算法都可以在GPU/TPU上运行,支持批处理和自动微分,为各种优化问题提供了灵活的解决方案。

使用示例

下面是一个使用JAXopt进行简单二次优化的示例:

import jax.numpy as jnp
from jaxopt import LBFGS

def objective(x):
    return jnp.sum((x - 1.0) ** 2)

solver = LBFGS(fun=objective)
x0 = jnp.zeros(5)
solution = solver.run(x0).params

print(f"Optimal solution: {solution}")

这个例子展示了如何定义一个目标函数,创建一个L-BFGS求解器,并使用它来找到最优解。

高级特性

隐式微分

JAXopt支持通过隐式微分来计算优化问题解对输入参数的梯度。这在双层优化和超参数优化等场景中非常有用。例如:

import jax

def f(x, theta):
    return jnp.sum((x - theta) ** 2)

def solve(theta):
    solver = LBFGS(lambda x: f(x, theta))
    return solver.run(jnp.zeros_like(theta)).params

dsolve = jax.grad(lambda theta: jnp.sum(solve(theta)))
print(dsolve(jnp.array([1.0, 2.0, 3.0])))

批处理优化

利用JAX的vmap功能,JAXopt可以轻松地并行处理多个优化问题:

from jax import vmap

def f(x, y):
    return (x - 2) ** 2 + (y - 3) ** 2

solver = LBFGS(f)
initial_guesses = jnp.array([[0., 0.], [1., 1.], [2., 2.]])
vmap_solve = vmap(solver.run)
solutions = vmap_solve(initial_guesses).params

print(f"Batch solutions: {solutions}")