Flax框架：Flax模块系统nn.Module

可视化演示

展示nn.Module生命周期（setup/call）、compact模式与setup模式对比、以及嵌套模块参数树结构，揭示Flax模型开发核心抽象。

模块基类：从一张空白菜谱开始

每家餐厅都会为厨师准备一张空白菜谱模板：上方是“食材清单”的空白栏，下方是“烹饪步骤”的空白栏。在 Flax 的世界里，nn.Module 就是这张模板。你不需要自己购买厨具、管理冰箱，模板自带了一套完整的“食材登记系统”——你只需回答两个问题：这道菜需要哪些食材？按什么顺序做？剩下的交给框架。

在代码里，继承 nn.Module 就能得到这套能力。setup 方法用来声明参数，它相当于填写食材清单；__call__ 方法用来定义前向计算，相当于写烹饪步骤。这两个方法的分工是 Flax 的核心设计：参数定义与计算逻辑分离，让模型结构一目了然。以最简单的线性层为例，它只有一个权重和一个偏置，前向计算就是矩阵乘法加偏置：

$$\displaystyle y = W x + b$$

1. 输入向量 x：代表进入这道菜的原料。在神经网络里，它就是上一层传下来的特征向量，长度为输入维度。

$$x \in \mathbb{R}^{d_{\text{in}}}$$

2. 权重矩阵 W：代表这道菜的“调味配方”，是模型要学习的核心参数之一，维度为 d_out × d_in，把输入从 d_in 维变换到 d_out 维。

$$W \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$$

3. 偏置向量 b：代表“基础底味”。即使原料全部为零，输出也带一个基准值，维度与输出一致。

$$b \in \mathbb{R}^{d_{\text{out}}}$$

4. 线性变换 Wx + b：把原料按配方混合并加上底味，得到最终成品，展开写成逐元素形式，更容易看清每个输出分量如何由所有输入加权求和得到。

$$y_i = \sum_{j=1}^{d_{\text{in}}} W_{ij} x_j + b_i$$

在 Flax 中，setup 里用 self.param 把 W 登记进食材清单，__call__ 里执行上述公式。套用这张菜谱模板，再复杂的模型也不过是“食材清单越来越长、烹饪步骤越来越多”而已。

compact 模式：快手厨师的边做边加料

有的厨师喜欢提前写备料单，有的厨师则习惯边炒菜边从冰箱里拿料——味道照样好，还省了写单子的时间。Flax 的 @nn.compact 装饰器就是为后者准备的“快手菜谱”模式。它允许你在 __call__ 方法内部直接内联声明子层，参数会像魔术一样自动被收集起来。

普通模式需要先把“食材”写进 setup，比如 self.dense = nn.Dense(8)。而在 compact 模式下，你甚至不需要 setup 方法，直接在 __call__ 里写 x = nn.Dense(8)(x)。这一行看似普通的调用，实际上同时完成了“声明一个 Dense 层”和“调用它”两件事。Flax 在第一次执行时会识别出这是一个子模块，自动分配名称（Dense_0、Dense_1……），并把它声明的参数挂到参数树上。下面这个两层 MLP 就是 compact 模式的典型例子：

$$\displaystyle f_{\theta}(x) = \text{Dense}_1\left( \text{ReLU}\left( \text{Dense}_0(x) \right) \right)$$

1. 内联层 Dense_0：第一个全连接层在 __call__ 内部直接写出，无需先在 setup 中赋值，它负责把输入升维到隐藏维度。

$$\text{Dense}_0: \mathbb{R}^{d_{\text{in}}} \to \mathbb{R}^{d_{\text{hidden}}}$$

2. 激活函数 ReLU：对中间结果做非线性变换，给模型引入学习复杂模式的能力，相当于给菜加一道“提鲜”步骤。

$$\text{ReLU}(z) = \max(0, z)$$

3. 内联层 Dense_1：第二个全连接层同样内联声明，把隐藏向量压缩到输出维度。

$$\text{Dense}_1: \mathbb{R}^{d_{\text{hidden}}} \to \mathbb{R}^{1}$$

4. 复合前向 f_θ：整个网络是两层变换的复合函数，参数 θ 由 compact 自动收集，无需手工拼装。

$$\theta = \left\{ \theta_{\text{Dense}_0}, \theta_{\text{Dense}_1} \right\}$$

compact 模式的最大价值在于代码精简：模型结构直接写在计算逻辑里，读代码就像读菜谱步骤一样流畅。当然，“边做边加料”也有代价——如果你的模型有大量需要复用的子模块，还是建议用 setup 显式声明，两者各有所长。

模块嵌套：套餐里的子菜谱

一张菜谱只能做一道菜，但餐厅的招牌套餐往往包含主菜、配菜和汤，每道菜都有自己独立的菜谱，组合在一起才构成完整的套餐。Flax 的模块嵌套正是这种哲学：一个 nn.Module 的内部可以持有其他 nn.Module 实例，小模块像积木一样拼成大模型。

在父模块的 setup 中，只需要写上 self.dense = DenseLayer(8)、self.mlp = MLP(16) 这样的赋值语句，Flax 就会在初始化时自动识别这些子模块，并递归地把它们的参数收集起来，构建出一棵参数树。你完全不需要手动把子模块的参数 dict 拼接起来——这就像套餐里的每道菜都有自己的食材清单，而后厨总管会自动把所有清单汇总进总账本。

$$\displaystyle \theta = \Big\{ \text{dense}: \{w, b\}, \quad \text{mlp}: \big\{ \text{Dense}_0: \{w_0, b_0\}, \text{Dense}_1: \{w_1, b_1\} \big\} \Big\}$$

1. dense 子模块：父模块的第一个子模块，它自己的参数 w 和 b 被挂在 dense 前缀下，形成参数树的第一层分支。

$$w \in \mathbb{R}^{4 \times 8}, \quad b \in \mathbb{R}^{8}$$

2. mlp 子模块：第二个子模块，而它内部又包含两个 Dense 层，形成嵌套的第二层分支——子模块的参数继续按它自己的内部结构递归展开。

$$\theta_{\text{mlp}} = \{ \theta_{\text{Dense}_0}, \theta_{\text{Dense}_1} \}$$

3. 树状展开：整个参数 θ 由递归规则展开唯一确定，父模块只需知道子模块是谁，细节全部自动完成。

$$\text{Leaves}(\theta) = \bigcup_{c \in \text{children}(\theta)} \text{Leaves}(\theta_c)$$

嵌套带来的另一个好处是模块复用。写好的 DenseLayer、MLP 就像封装好的菜谱，可以出现在任意父模块里，既保证风格统一，又避免重复劳动。这也是为什么 Flax 能优雅地组织从几层小网络到上千层大模型的结构。

参数树结构：锁进保险柜的总账本

model.init(key, x) 是 Flax 中“开张备货”的时刻：它根据输入的形状，按模块声明生成所有参数，并返回一棵完整的参数树。这棵树不是普通的 Python 字典，而是一个只读的 FrozenDict——就像锁进保险柜的后厨总账本，你可以翻看每一页，但不能直接涂改。

FrozenDict 之所以只读，背后是 JAX 的函数式哲学：参数是纯数据，任何“修改”都应生成新对象，而不是就地改动。这保证了训练过程可复现、可并行。那么想更新参数怎么办？答案是先解冻（unfreeze），改成普通 dict，更新完毕后再冻结（freeze）。在 Flax 里，每次梯度更新本质上都是“复印账本 → 修正数字 → 重新锁柜”的过程。

$$\displaystyle \text{tree\_flatten}(\theta) = \Big( [w, b, w_0, b_0, w_1, b_1], \; \text{treedef} \Big)$$

1. 叶子列表：把参数树最底层的每个具体数组按深度优先顺序摊平，得到一维列表，其中每个元素都是 JAX 数组。

$$\text{leaves} = [w, b, w_0, b_0, w_1, b_1]$$

2. treedef 结构描述：记录每个叶子在原树中的嵌套路径（如 dense/w、mlp/Dense_0/kernel），有了它就能把叶子列表重新组装回原树。

$$\text{treedef}: \text{dense}\{w, b\}, \text{mlp}\{\text{Dense}_0\{w_0, b_0\}, \text{Dense}_1\{w_1, b_1\}\}$$

3. 只读约束：FrozenDict 直接赋值会抛异常，必须先 unfreeze 成普通 dict 才能改写，再 freeze 回去。

$$\theta' = \text{freeze}\left( \text{modify}\left( \text{unfreeze}(\theta) \right) \right)$$

把参数树设计成冻结结构，看似增加了“解冻—冻结”的步骤，实则换来了极大的安全性：任何函数拿到参数树都只能读取、不能破坏，这让 JAX 的 jit 编译和自动微分能够放心大胆地优化整个计算图。

树工具遍历：后厨总管的一把算盘

账本有了，还得有人会查账。jax.tree_util 就是那位拿着总账本的后厨总管，它提供了一套操作参数树的通用工具。其中最常用的是 tree_map：它把一个函数应用到参数树的每一个叶子上，并保持树的结构不变。比如想给所有参数统一放大 0.1 倍，或者同时给每个参数加上一个对应梯度，tree_map 都能一行搞定。

tree_map 真正的威力体现在训练中。当你用 jax.grad 算出梯度后，梯度 grads 本身也是一棵与 params 结构完全相同的树。把两者一起喂给 tree_map，它就会自动把每个参数和它的梯度逐叶配对，执行参数更新：

$$\displaystyle g_t = \nabla_\theta \mathcal{L}(\theta_t), \qquad \theta_{t+1} = \text{tree\_map}\left( (p, g) \mapsto p - \eta g, \; \theta_t, \; g_t \right)$$

1. 梯度 g_t：jax.grad 对损失函数关于全部参数求导，得到与参数树同构的梯度树，每个梯度叶子和参数叶子一一配对。

$$g_t = \frac{\partial \mathcal{L}}{\partial \theta_t}$$

2. 学习率 η：控制每次更新迈多大的步子，是训练中最关键的超参数之一。

$$\eta \in (0, 1)$$

3. tree_map 配对更新：遍历两棵树的叶子，对每一对 (p, g) 执行 p - ηg，最终返回一棵更新后的新参数树，结构与原来完全一致。

$$\theta_{t+1} = \theta_t - \eta g_t$$

除了 tree_map，tree_leaves 可以快速取出所有叶子（比如统计参数量），tree_flatten / tree_unflatten 负责摊平与还原。有了这套工具，无论是百万参数的小模型还是千亿参数的大模型，更新逻辑都只需要几行代码——后厨总管一把算盘打天下。

训练实践：试菜调出好味道

菜谱定稿前，总要反复试菜：先按菜谱做一桌，尝一尝，觉得咸了淡了，就在配方上略作调整，下一轮再试。Flax 模型的训练就是这样一个循环。整个流程可以拆成四个环节：init（开张备好初始配方）、apply（按菜谱做一桌菜）、grad（试吃找出偏差）、tree_map（调整配方再试）。

初始化的目标是生成参数树，model.init(key, x) 会根据输入形状推断每个参数的尺寸，并使用各种初始化策略设置初始值——这相当于厨师根据食材备好第一版配方。前向计算则由 model.apply(params, x) 完成：它拿着参数树，按模块的 __call__ 逻辑把数据从头传到尾，得到预测值。

$$\displaystyle \theta^* = \arg\min_{\theta} \; \frac{1}{N} \sum_{i=1}^{N} \mathcal{L}\left( f_\theta(x_i), y_i \right)$$

1. 训练样本对：数据集中一对对输入 x_i 与标签 y_i，是试菜的“标准口味”。

$$(x_i, y_i), \quad i = 1, 2, \ldots, N$$

2. 损失函数 L：衡量预测值与真实标签的差距，均方误差、交叉熵都是 L 的常见形式，公式给出均方误差的写法。

$$\mathcal{L} = \frac{1}{N} \sum_{i=1}^{N} \left( f_\theta(x_i) - y_i \right)^2$$

3. 梯度下降寻优：通过计算损失对参数的梯度并沿负方向更新，反复迭代逼近最优配方，正是上一节 tree_map 的用武之地。

$$\theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L}(\theta_t)$$

把四个环节串起来，就是一个标准的 Flax 训练步：grads = jax.grad(loss_fn)(params) 算出梯度，jax.tree_util.tree_map 完成更新。这一套流程从 MNIST 手写识别到 GPT 级大模型从未改变——菜谱可以越来越复杂，但“开张、做菜、试吃、调配方”的循环永远不变。

AI 应用场景

Python 演示

import jaximport jax.numpy as jnpfrom flax import linen as nnfrom jax import random  # 1) setup 风格: 先列食材清单, 再写烹饪步骤class DenseLayer(nn.Module):    out_dim: int     def setup(self):        self.w = self.param(            'w', nn.initializers.lecun_normal(), (4, self.out_dim))        self.b = self.param(            'b', nn.initializers.zeros, (self.out_dim,))     def __call__(self, x):        return x @ self.w + self.b  # 2) compact 风格: 边做菜边加料, 一行声明一个层class MLP(nn.Module):    hidden: int     @nn.compact    def __call__(self, x):        x = nn.Dense(self.hidden)(x)        x = jax.nn.relu(x)        x = nn.Dense(1)(x)        return x  # 3) 模块嵌套: 小模块拼成套餐, 参数树自动构建class Net(nn.Module):     def setup(self):        self.dense = DenseLayer(8)        self.mlp = MLP(16)     def __call__(self, x):        return self.mlp(self.dense(x))  def main():    key = random.PRNGKey(42)    x = jnp.ones((5, 4))    y = jnp.ones((5, 1))    model = Net()    params = model.init(key, x)    print('auto-built parameter tree:')    print(params)     # 前向传播: 按菜谱把食材煮成成品    pred = model.apply(params, x)    print('forward output shape:', pred.shape)     def loss_fn(p):        return jnp.mean((model.apply(p, x) - y) ** 2)     # 试吃: 用自动微分算出每个参数的"偏差方向"    grads = jax.grad(loss_fn)(params)     def update(p, g):        return p - 0.01 * g     # 后厨总管 tree_map: 逐叶配对参数与梯度, 一次完成更新    new_params = jax.tree_util.tree_map(update, params, grads)    print('first weight shape after update:',          jax.tree_util.tree_leaves(new_params)[0].shape)  if __name__ == '__main__':    main()

1. self.param('w', nn.initializers.lecun_normal(), (4, self.out_dim))在 setup 中声明一个参数并指定初始化器与形状，相当于把“食材 w”记入菜谱清单。
2. @nn.compact装饰 call，开启“边做菜边加料”模式，让内联的 nn.Dense 自动注册参数。
3. x = nn.Dense(self.hidden)(x)compact 模式下“声明即调用”：一行代码同时完成子层实例化与前向计算。
4. model.init(key, x)根据输入形状生成整棵 FrozenDict 参数树，即“开张备货”。
5. model.apply(params, x)拿着参数树执行前向传播，对应“按菜谱做一桌菜”。
6. jax.grad(loss_fn)(params)对损失函数求关于所有参数的梯度，得到的梯度树与参数树结构一致。
7. jax.tree_util.tree_map(update, params, grads)后厨总管逐叶配对参数与梯度，一行完成全量参数更新。

常见误区

面试题