DiffEqFlux.jl

DiffEq(For)Lux.jl（又称 DiffEqFlux.jl）通过帮助用户将微分方程求解器嵌入到神经网络中，实现了微分方程领域与机器学习领域的融合。该包以 DifferentialEquations.jl 和 Lux.jl 为基石，支持 科学机器学习 领域的研究，尤其是神经微分方程，从而将物理信息融入传统机器学习模型中。

[!注意] 我们通过 FromFluxAdaptor() 保持与 Flux.jl 的向后兼容性。

教程与文档

有关如何使用本包的信息，请参阅 稳定版文档。如需查看包含未发布功能的版本文档，请使用 开发中文档。

问题领域

DiffEqFlux.jl 适用于隐式层机器学习。它提供了与 Flux.jl 和 Lux.jl 等机器学习库接口兼容的架构，便于将连续时间机器学习层集成到更大的机器学习应用中。

现有以下层函数：

• 神经常微分方程（Neural ODEs）
• 基于配点法的神经常微分方程（无需求解器，速度最快的方式！）
• 多重打靶神经常微分方程
• 神经随机微分方程（Neural SDEs）
• 神经微分代数方程（Neural DAEs）
• 神经时滞微分方程（Neural DDEs）
• 扩展型神经常微分方程
• 哈密顿神经网络（配备专门的二阶和辛积分器）
• 连续归一化流（CNF）和 FFJORD

这些方法支持高阶、自适应、隐式、GPU 加速、牛顿-克雷洛夫等求解技术。更多示例请参考 发布博客文章。此外，关于神经偏微分方程和神经跳跃随机微分方程等更多演示，可参见 这篇博客文章（以及其他许多内容）。

不要局限于当前的神经化方式。借助本包，您可以探索多种整合这两种方法论的方式：

• 可以定义“激活”由微分方程描述的非线性函数的神经网络。
• 可以定义部分层为常微分方程求解的神经网络。
• 可以定义部分项为神经网络的常微分方程。
• 可以在常微分方程上定义成本函数来生成神经网络。

破坏性变更

v4

• TensorLayer 已被移除，改用 Boltz.Layers.TensorProductLayer。
• DiffEqFlux 中的基础函数已被移除，取而代之的是 Boltz.Basis 模块。
• SplineLayer 已被移除，改用 Boltz.Layers.SplineLayer。
• NeuralHamiltonianDE 已被移除，改用带有 Layers.HamiltonianNN 的 NeuralODE。
• HamiltonianNN 已被移除，取而代之的是 Layers.HamiltonianNN。
• Lux 和 Boltz 更新至 v1 版本。

v3

• 移除了对 Flux 的依赖。如果传入非 Lux 的 AbstractLuxLayer，我们会尝试使用 FromFluxAdaptor()(model) 自动将其转换为 Lux 模型。
• Flux 不再从 DiffEqFlux 中重新导出，而是改由 Lux 重新导出。
• NeuralDAE 现在允许将可选的 du0 作为输入。
• TensorLayer 现已成为一个 Lux 神经网络。
• 多个层构造的 API 发生了变化。详情请参阅更新后的文档。

DiffEqFlux.jl 快速上手指南

DiffEqFlux.jl 是一个将微分方程求解器与机器学习深度融合的 Julia 库。它允许用户在神经网络中嵌入微分方程求解层，从而实现神经微分方程（Neural ODEs/SDEs/DAEs 等），特别适用于科学机器学习（Scientific ML）领域，能够将物理规律融入传统机器学习模型。

环境准备

• 操作系统：支持 Linux、macOS 和 Windows。
• Julia 版本：建议安装最新稳定版 Julia（1.9 或更高版本）。
• 前置知识：熟悉基本的 Julia 语法，了解 Lux.jl（现代神经网络库）和 DifferentialEquations.jl（微分方程求解库）的基本概念会有所帮助。
• 依赖说明：本工具核心基于 Lux.jl 构建。虽然不再直接依赖 Flux.jl，但提供了适配器以兼容旧的 Flux 模型。

安装步骤

在 Julia REPL 或脚本中，进入包管理模式并安装 DiffEqFlux：

using Pkg
Pkg.add("DiffEqFlux")

提示：国内开发者若遇到下载速度慢的问题，可配置国内镜像源加速安装。在运行上述命令前，执行以下设置（以清华源为例）：

ENV["JULIA_PKG_SERVER"] = "https://mirrors.tuna.tsinghua.edu.cn/julia"

安装完成后，加载包即可使用：

using DiffEqFlux
using Lux, Random

基本使用

以下示例展示如何构建并训练一个最简单的**神经常微分方程（Neural ODE）**模型。该模型将神经网络作为微分方程的动态项进行优化。

1. 定义问题与数据

假设我们要拟合一个简单的正弦波数据 $u' = f(u, p, t)$。

using OrdinaryDiffEq, Plots

# 生成训练数据 (真值: u' = -0.5u + cos(t))
tspan = (0.0f0, 10.0f0)
tsteps = 0.0f0:0.1f0:10.0f0
u0 = [2.0f0]

# 真实动力学函数
function true_dudt(u, p, t)
    return [-0.5u[1] + cos(t)]
end

prob_true = ODEProblem(true_dudt, u0, tspan)
sol_true = solve(prob_true, Tsit5(), saveat=tsteps)
X_train = sol_true.u
Y_train = sol_true.u # 自监督学习，预测自身轨迹

2. 构建神经 ODE 层

使用 Lux 定义神经网络，并将其包裹在 NeuralODE 中。

# 定义一个简单的全连接神经网络 (输入 1 -> 隐藏层 32 -> 输出 1)
nn = Lux.Chain(
    Lux.Dense(1, 32, tanh),
    Lux.Dense(32, 1)
)

# 初始化参数
rng = Random.default_rng()
p_init, st = Lux.setup(rng, nn)

# 构建 NeuralODE 对象
# 使用 Tsit5() 求解器，自动处理反向传播
neural_ode = NeuralODE(nn, tspan, Tsit5(), saveat=tsteps)

# 包装成可训练的模型函数
function predict(p)
    Array(neural_ode(u0, p, st)[1])
end

3. 定义损失函数并训练

使用 Optimisers.jl 或 SciMLBase 提供的训练循环进行参数优化。

using Optimisers

# 损失函数：预测轨迹与真实轨迹的均方误差
function loss(p)
    pred = predict(p)
    sum(abs2, X_train .- pred)
end

# 初始损失
println("Initial Loss: $(loss(p_init))")

# 设置优化器 (Adam)
opt = Optimisers.Adam(0.05)
opt_state = Optimisers.setup(opt, p_init)

# 简单训练循环 (实际应用中建议使用 SciML 的训练回调)
for i in 1:100
    grad = Flux.gradient(loss, p_init) # 注意：此处需确保梯度追踪正确，推荐使用 Zygote 或 Enzyme
    # 由于 DiffEqFlux 深度集成，通常直接使用 sciml 训练接口更稳健
    # 这里仅展示逻辑，具体梯度计算依赖后端配置
    println("Step $i completed") 
end

注意：在实际生产代码中，推荐使用 SciMLBase 生态提供的标准训练接口（如 sciml_train）来处理微分方程反向传播的复杂性，以获得最佳的 GPU 加速和稳定性支持。

通过以上步骤，您已成功构建了一个包含微分方程求解层的神经网络。您可以进一步探索神经 SDE、神经 DAE 或哈密顿神经网络等高级架构。详细 API 请参考官方文档。