Flax NNX 基础教程：JAX 上的神经网络构建新范式

Flax NNX 基础教程：JAX 上的神经网络构建新范式

flax Flax is a neural network library for JAX that is designed for flexibility. 项目地址: https://gitcode.com/gh_mirrors/fl/flax

前言

在深度学习框架领域，JAX 因其函数式编程特性和强大的自动微分能力而备受关注。然而，JAX 的纯函数式特性也给模型构建带来了挑战。Flax NNX 作为 Flax 生态系统的新成员，通过引入 Python 原生对象语义，为 JAX 带来了更直观的神经网络构建体验。

Flax NNX 模块系统

模块设计理念

Flax NNX 的 nnx.Module 系统采用显式设计原则，与传统的 Flax Linen 或 Haiku 有着显著区别：

1. 状态显式持有：模块直接维护自身的参数状态
2. PRNG 状态显式传递：用户需要显式管理随机数生成器状态
3. 形状显式指定：初始化时必须提供完整的形状信息，不进行形状推断

线性层实现示例

class Linear(nnx.Module):
  def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
    key = rngs.params()
    self.w = nnx.Param(jax.random.uniform(key, (din, dout)))
    self.b = nnx.Param(jnp.zeros((dout,)))
    self.din, self.dout = din, dout

  def __call__(self, x: jax.Array):
    return x @ self.w + self.b

关键点说明：

• nnx.Param 封装可训练参数
• 静态属性（如维度信息）直接存储
• __call__ 方法实现前向计算

状态可视化

Flax NNX 提供了便捷的模型可视化工具：

model = Linear(2, 5, rngs=nnx.Rngs(params=0))
y = model(x=jnp.ones((1, 2)))
nnx.display(model)

可视化基于 Treescope 库实现，可以清晰展示模型结构和参数。

状态管理与模块组合

状态更新机制

Flax NNX 支持在正向传播过程中更新状态，这对于实现批归一化等层非常关键：

class Counter(nnx.Module):
  def __init__(self):
    self.count = Count(jnp.array(0))

  def __call__(self):
    self.count += 1  # 直接修改状态

模块嵌套

模块可以自由组合形成复杂结构：

class MLP(nnx.Module):
  def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):
    self.linear1 = Linear(din, dmid, rngs=rngs)
    self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)
    self.bn = nnx.BatchNorm(dmid, rngs=rngs)
    self.linear2 = Linear(dmid, dout, rngs=rngs)

模型手术与参数共享

Flax NNX 的模块具有高度灵活性：

# 将普通线性层替换为LoRA层
model.linear1 = LoraLinear(model.linear1, 4, rngs=rngs)

这种设计使得模型结构调整、参数共享等操作变得非常简单。

Flax 转换系统

训练步骤实现

@nnx.jit
def train_step(model, optimizer, x, y):
  def loss_fn(model: MLP):
    y_pred = model(x)
    return jnp.mean((y_pred - y) ** 2)

  loss, grads = nnx.value_and_grad(loss_fn)(model)
  optimizer.update(grads)  # 原地更新参数
  return loss

关键特性：

• 状态更新自动传播
• 优化器与模型保持引用关系
• 支持 JIT 编译

层堆叠与扫描

@nnx.scan(in_axes=(0, nnx.Carry), out_axes=nnx.Carry)
def forward(model: MLP, x):
  x = model(x)
  return x

nnx.scan 提供了比 jax.lax.scan 更直观的接口，支持多输入输出和扫描轴指定。

Flax 函数式 API

状态与图定义分离

graphdef, state = nnx.split(model)

• State: 包含所有可变状态的映射
• GraphDef: 包含重建模块所需的静态信息

合并与更新

# 重建模块
model = nnx.merge(graphdef, state)

# 更新模块状态
nnx.update(model, new_state)

这种分离设计使得模块可以无缝集成到 JAX 的转换系统中。

总结

Flax NNX 通过以下特性显著改善了 JAX 上的深度学习体验：

1. 直观的对象语义：类似 PyTorch 的编程模型
2. 灵活的状态管理：支持训练过程中的状态更新
3. 强大的组合性：模块可以自由嵌套和重组
4. 无缝 JAX 集成：保持与 JAX 生态的兼容性

对于从 PyTorch 或 Keras 迁移的用户，Flax NNX 提供了熟悉而强大的接口，同时保留了 JAX 的性能优势。

flax Flax is a neural network library for JAX that is designed for flexibility. 项目地址: https://gitcode.com/gh_mirrors/fl/flax