Haiku
DeepMind为JAX写的面向对象神经网络库

🎯什么是Haiku：Sonnet for JAX

在JAX的世界里，一切都是纯函数。jax.grad求导、jax.jit编译加速、jax.pmap并行化——这些强大的变换都要求输入函数没有副作用、没有可变状态、没有隐式参数。这对于写一个矩阵乘法来说很自然，但对于构建一个有数百层、数千个参数、需要随机采样的深度神经网络来说，纯函数式的编程模型会变得极其繁琐。

DeepMind的答案是Haiku。正如项目首页那句俳句：

Haiku is a tool
For building neural networks
Think: "Sonnet for JAX"

Sonnet是DeepMind在TensorFlow时代广泛使用的神经网络库，由DeepMind的研究者团队开发，在内部几乎被所有项目采用。Sonnet的核心编程模型是Module——一个Python对象，持有自己的参数、子模块和方法。这种面向对象的模型让研究者可以用熟悉的class写法定义网络，而不用手写参数字典的嵌套结构。

Haiku的作者正是Sonnet的部分原作者。他们将Sonnet的Module编程模型移植到JAX上，但做了一个关键的架构决策：Module本身是"不纯"的（它们在内部调用hk.get_parameter来隐式管理参数），但通过hk.transform，这些不纯的Module被变换成JAX可以消费的纯函数。这就是Haiku名字的由来——俳句（Haiku）是一种日本诗歌形式，而Sonnet是一种英国诗歌形式；从Sonnet到Haiku，是从TensorFlow到JAX的隐喻。

🏗️核心设计：hk.transform的纯函数变换

Haiku的全部魔法浓缩在一个函数里：hk.transform。这个函数接受一个"不纯"的函数（内部使用了hk.Module、hk.get_parameter等Haiku原语），返回一对纯函数：init和apply。

让我们用Haiku官方的MNIST示例来理解这个变换的完整流程：

import haiku as hk
import jax
import jax.numpy as jnp

def loss_fn(images, labels):
    mlp = hk.Sequential([
        hk.Linear(300), jax.nn.relu,
        hk.Linear(100), jax.nn.relu,
        hk.Linear(10),
    ])
    logits = mlp(images)
    return jnp.mean(softmax_cross_entropy(logits, labels))

# 核心变换：不纯函数 → 纯函数对
loss_fn_t = hk.transform(loss_fn)
loss_fn_t = hk.without_apply_rng(loss_fn_t)

rng = jax.random.PRNGKey(42)
dummy_images, dummy_labels = next(input_dataset)

# init：运行函数，收集所有参数
params = loss_fn_t.init(rng, dummy_images, dummy_labels)

# apply：注入参数，执行前向计算
loss = loss_fn_t.apply(params, None, images, labels)

# 因为apply是纯函数，可以直接传给jax.grad
grads = jax.grad(loss_fn_t.apply)(params, images, labels)

这段代码揭示了Haiku的核心工作流程：

init阶段：Haiku运行你的函数，但不是真的执行计算——它是在"追踪"（trace）执行过程。每当hk.get_parameter被调用（比如在hk.Linear的内部），Haiku会记录这个参数的名称、形状、 dtype和初始化器，然后实际调用初始化器生成初始值。函数运行结束后，所有被收集的参数以嵌套字典的形式返回。

返回的params结构长这样：

{
    'linear': {
        'b': ndarray(..., shape=(300,), dtype=float32),
        'w': ndarray(..., shape=(28, 300), dtype=float32)
    },
    'linear_1': {
        'b': ndarray(..., shape=(100,), dtype=float32),
        'w': ndarray(..., shape=(1000, 100), dtype=float32)
    },
    'linear_2': {
        'b': ndarray(..., shape=(10,), dtype=float32),
        'w': ndarray(..., shape=(100, 10), dtype=float32)
    }
}

apply阶段：Haiku再次运行你的函数，但这次每当hk.get_parameter被调用时，它不再创建新参数，而是从你传入的params字典中查找对应的值。这意味着同一个函数定义可以被复用——init时用于初始化参数，apply时用于前向推理，而apply是纯函数，可以安全地传给jax.grad、jax.jit等变换。

hk.without_apply_rng是一个便利包装器。如果你的前向计算不需要随机数（比如纯确定性的MLP），你可以用这个函数移除apply的rng参数，使调用更简洁。如果需要随机数（比如Dropout、VAE），则保留rng参数。

🧩hk.Module：面向对象的参数管理

Haiku的Module是Sonnet Module的直接对应物。所有模块都是hk.Module的子类，通常实现__init__和__call__方法：

class MyLinear(hk.Module):

    def __init__(self, output_size, name=None):
        super().__init__(name=name)
        self.output_size = output_size

    def __call__(self, x):
        j, k = x.shape[-1], self.output_size
        w_init = hk.initializers.TruncatedNormal(1. / np.sqrt(j))
        w = hk.get_parameter("w", shape=[j, k], dtype=x.dtype, init=w_init)
        b = hk.get_parameter("b", shape=[k], dtype=x.dtype, init=jnp.zeros)
        return jnp.dot(x, w) + b

这里有几个关键设计点值得深入理解：

命名自动推导：每个Module都有一个名字。如果你不传name参数，Haiku会从类名自动推导——MyLinear变成my_linear。当同一个函数中实例化多个同类型模块时，Haiku会自动添加后缀（my_linear、my_linear_1、my_linear_2...）。这个名字构成了params字典的键，确保参数的唯一可寻址性。

hk.get_parameter而非对象属性：注意参数不是通过self.w = ...存储的，而是通过hk.get_parameter("w", ...)获取的。这是Haiku设计的核心——Module对象本身不持有参数值，参数值由Haiku的内部上下文管理。init时上下文创建并收集参数，apply时上下文从外部注入参数。这就是为什么同一个Module类可以在init和apply中复用——对象本身是无状态的，状态由外部的params字典承载。

嵌套模块：Module可以包含其他Module，形成树状结构。参数字典的嵌套结构会自动反映模块的嵌套关系。例如，一个自定义模块包含一个hk.nets.MLP和一个自定义Linear，params结构会是：

{
    'custom_linear/~/hk_internal_linear/~/linear_0': {'b': ..., 'w': ...},
    'custom_linear/~/hk_internal_linear/~/linear_1': {'b': ..., 'w': ...},
    'custom_linear/~/old_linear': {'b': ..., 'w': ...}
}

其中/~/是Haiku的模块层级分隔符，类似于文件系统路径。这种设计让你可以清晰地定位和操作任意层级的参数。

🎲RNG密钥管理：hk.next_rng_key

JAX使用显式的PRNG（伪随机数生成器）密钥——每次随机操作都需要传入一个jax.random.PRNGKey。在纯JAX中，如果你需要多个随机数，你必须手动split密钥：

key1, key2 = jax.random.split(key)
a = jax.random.normal(key1, shape=(100,))
b = jax.random.bernoulli(key2, 0.5, shape=(100,))

在一个大型网络中手动管理密钥分裂是极其痛苦的——你需要为Dropout、VAE采样、数据增强等每一个随机操作都预留一个独立的密钥。Haiku用一行代码解决了这个问题：

class MyDropout(hk.Module):

    def __init__(self, rate=0.5, name=None):
        super().__init__(name=name)
        self.rate = rate

    def __call__(self, x):
        key = hk.next_rng_key()
        p = jax.random.bernoulli(key, 1.0 - self.rate, shape=x.shape)
        return x * p / (1.0 - self.rate)

hk.next_rng_key()在transformed函数内部每次调用都返回一个唯一的、确定性的RNG密钥。这些密钥从一个初始密钥（传入init或apply的rng参数）确定性派生，因此：

同一个初始密钥 + 同一个函数调用顺序 = 完全相同的随机序列。这保证了JAX程序变换（尤其是jax.jit和jax.pmap）的正确性——无论函数被trace多少次，随机行为都是可复现的。

Haiku内部维护一个RNG密钥序列，每次hk.next_rng_key()调用都会从序列中取下一个密钥并分裂出新的序列密钥。这个过程对用户完全透明——你只需要在顶层传入一个PRNGKey，Haiku自动处理所有的分裂和派生。

📦有状态模块：hk.transform_with_state

有些模块需要维护非训练参数的内部状态——最典型的例子是Batch Normalization中的移动平均值。在训练过程中，BN层需要维护输入的移动平均统计量，用于推理时的归一化。这些统计量不是通过梯度下降学习的，但需要在每次前向传播时更新。

Haiku通过hk.transform_with_state处理这种情况。与hk.transform不同，它返回的init和apply函数有额外的state参数：

def forward(x, is_training):
    net = hk.nets.ResNet50(1000)
    return net(x, is_training)

forward = hk.transform_with_state(forward)

# init返回params和state
params, state = forward.init(rng, x, is_training=True)

# apply接收params和state，返回输出和更新后的state
logits, state = forward.apply(params, state, rng, x, is_training=True)

state的结构与params相同——一个按模块名嵌套的字典。在ResNet50的例子中，state会包含所有BatchNorm层的移动平均值和移动方差。每次apply调用后，state被更新并返回，你需要在下一次调用时传入更新后的state。

Haiku提供了hk.get_state和hk.set_state两个原语来读写状态：

def stateful_f(x):
    counter = hk.get_state("counter", shape=[], dtype=jnp.int32, init=jnp.ones)
    multiplier = hk.get_parameter('multiplier', shape=[1,], dtype=x.dtype, init=jnp.ones)
    hk.set_state("counter", counter + 1)
    output = x + multiplier * counter
    return output

stateful_forward = hk.without_apply_rng(hk.transform_with_state(stateful_f))

如果你忘记使用hk.transform_with_state而代码中调用了hk.set_state，Haiku会打印清晰的错误信息指向正确的变换函数，而不是静默丢弃你的状态。

⚡分布式训练：与jax.pmap无缝协作

Haiku返回的纯函数与jax.pmap完全兼容，这使得数据并行训练变得直接。核心模式是：

# 初始化模型（单设备）
rng = jax.random.PRNGKey(428)
params = loss_fn_t.init(rng, sample_image, sample_label)

# 复制参数到所有设备
num_devices = jax.local_device_count()
params = jax.tree.map(lambda x: np.stack([x] * num_devices), params)

def update(params, inputs, labels, axis_name='i'):
    grads = jax.grad(loss_fn_t.apply)(params, inputs, labels)
    # 跨设备梯度平均
    grads = jax.lax.pmean(grads, axis_name)
    new_params = my_update_rule(params, grads)
    return new_params

# 每个设备处理一个batch，梯度自动跨设备平均
for _ in range(10):
    superbatch_images, superbatch_labels = make_superbatch()
    params = jax.pmap(update, axis_name='i')(
        params, superbatch_images, superbatch_labels
    )

这里的关键是jax.lax.pmean——它对所有设备上的梯度取平均，确保所有设备使用相同的梯度更新参数，从而保持参数同步。Haiku不需要任何特殊处理就能支持这个模式，因为apply是纯函数——给定相同的params和输入，输出确定。

Haiku官方提供了完整的ResNet-50在ImageNet上的分布式训练示例，展示了从数据加载到多设备训练的完整流程。在DeepMind内部，Haiku被用于大规模的图像处理、语言模型和强化学习实验，这些实验通常在数百个TPU核心上运行。

🔬内置网络与嵌套模块

Haiku提供了一套完整的内置模块和网络，覆盖了常见的深度学习组件：

基础层：hk.Linear（全连接层）、hk.Conv2D（2D卷积）、hk.Conv1D、hk.Conv3D、hk.max_pool、hk.avg_pool等。这些层的API设计直接对应Sonnet，参数名、初始化默认值都与Sonnet保持一致。

完整网络：hk.nets.MLP（多层感知机）、hk.nets.ResNet50/101/152/200（残差网络）、hk.nets.MobileNetV1/V2、hk.nets.VAE（变分自编码器）。这些网络可以直接使用，也可以作为自定义模块的子组件。

模块组合：hk.Sequential按顺序组合多个模块，类似于PyTorch的nn.Sequential。嵌套模块的参数会自动按模块层级组织：

class MyModuleCustom(hk.Module):
    def __init__(self, output_size=2, name='custom_linear'):
        super().__init__(name=name)
        self._internal_linear_1 = hk.nets.MLP(
            output_sizes=[2, 3], name='hk_internal_linear'
        )
        self._internal_linear_2 = MyLinear1(
            output_size=output_size, name='old_linear'
        )

    def __call__(self, x):
        return self._internal_linear_2(self._internal_linear_1(x))

这种组合模式让复杂网络的构建变得直观——你用class写法定义网络结构，Haiku自动处理参数的收集、命名和嵌套。

其他工具：Haiku还提供了hk.dropout、hk.BatchApply（将前两维合并以处理变长序列）、hk.MultiHeadAttention（多头注意力机制）、hk.Embed（词嵌入）等常用组件。此外还有与Flax的互操作模块（haiku.experimental.flax），允许在Haiku中使用Flax模块或反之。

⚠️已知限制与JAX变换嵌套陷阱

Haiku的文档中有一个醒目的警告：在Haiku网络内部使用JAX变换（如jax.jit和jax.remat）可能导致难以解释的tracing错误，甚至静默产生错误结果。

问题的根源在于Haiku的上下文管理机制。hk.transform通过一个内部上下文（context）来追踪参数和RNG密钥。当你在Haiku函数内部使用jax.jit时，jit会重新trace你的函数，但此时Haiku的上下文可能不存在或处于不一致的状态。

Haiku为此提供了一组专门的包装器来安全地在Haiku网络内部使用JAX变换：

hk.vmap：在Haiku模块内部安全使用jax.vmap。例如，对一个模块的输出在某个维度上进行向量化映射。hk.scan：在Haiku模块内部安全使用jax.scan，常用于RNN的时间步展开。hk.remat：在Haiku模块内部安全使用jax.remat，用于梯度检查点以节省内存。

这些包装器的核心作用是在JAX变换的trace过程中正确地传播Haiku的上下文。如果你需要更细粒度的控制，Haiku还提供了hk.lift来将内层变换的参数提升到外层。

另一个需要注意的限制是参数共享。在标准的hk.transform中，每次在函数中实例化一个Module都会创建新的参数。如果你需要共享参数（比如Siamese网络），你需要使用hk.experimental.name_scope或在同一函数中重用同一个Module实例。

⚖️Haiku vs Flax：DeepMind的内部抉择

2023年7月，Google DeepMind发布了一条重要公告：推荐新项目采用Flax而非Haiku。这是一个具有深远影响的决策，值得深入分析其背后的原因。

Flax的功能超集：截至2023年，Flax已经具备了Haiku的所有核心功能——Module抽象、参数管理、状态处理、RNG管理——并且在某些方面走得更远。Flax的 linen.Module 支持更灵活的参数共享模式、更强大的条件分支（linen.cond），以及更完善的训练API（flax.training）。

社区与生态：Flax最初由Google Brain开发，后来与DeepMind合并后成为Google DeepMind的统一推荐。Flax拥有更大的开发团队、更活跃的社区、更丰富的文档和示例。Hugging Face的transformers库对Flax有一等支持，而Haiku的集成相对有限。

Haiku的维护模式：Haiku不会停止维护——DeepMind内部有大量Haiku代码在使用，他们会持续修复bug、保持与新版JAX和Python的兼容性。但不会再添加新功能，也不接受新功能的PR。这意味着Haiku是一个"完成"的项目——它的设计目标已经实现，代码库稳定。

迁移考量：对于已有Haiku项目，是否迁移到Flax取决于项目规模和活跃度。对于新的研究项目，直接使用Flax是更安全的选择。但理解Haiku的设计对于理解JAX的编程模型仍然有价值——hk.transform是理解"如何将面向对象代码嵌入纯函数世界"的优秀教材。

💡对研究者的启示

Haiku虽然进入维护模式，但它对JAX生态和深度学习研究方法论的影响是深远的。从Haiku的设计中，我们可以提炼出几个对研究者仍然有价值的启示。

面向对象与纯函数的桥梁。Haiku最核心的贡献是展示了如何在纯函数式编程（JAX）中保留面向对象编程的便利性。hk.transform的模式——用上下文追踪将OO代码变换为纯函数——是一种可复用的设计模式。无论你使用Haiku、Flax还是其他框架，理解这个变换过程有助于你写出更好的JAX代码。

参数管理的最佳实践。Haiku的params字典设计——按模块名嵌套、可inspect、可tree.map——已经成为JAX生态的事实标准。Flax的参数结构、Optax的优化器状态、Chex的数据结构验证都遵循类似的模式。学会阅读和操作这种嵌套字典结构是JAX开发者的基本技能。

RNG管理的简洁方案。hk.next_rng_key()的设计——从单一初始密钥确定性派生所有随机数——是JAX中处理随机性的优雅方案。即使你使用Flax（它有自己的RNG管理机制），理解这种"单一密钥 → 序列派生"的模式有助于你调试随机性相关的问题。

从Haiku到AlphaFold。Haiku在DeepMind内部被广泛使用，包括AlphaFold 2的早期开发。虽然AlphaFold 2的公开代码使用了JAX+Evoformer的定制实现，但DeepMind内部的许多原型实验是在Haiku中进行的。这意味着理解Haiku有助于你阅读和理解DeepMind的研究代码——许多论文的补充材料和开源代码仍然使用Haiku。

何时选择Haiku vs Flax。对于2026年的新项目，Flax是首选。但如果你需要阅读或复现DeepMind 2020-2023年间的开源代码（如某些RL算法、生成模型），Haiku知识仍然必要。此外，Haiku的简洁性（核心API只有hk.transform、hk.Module、hk.get_parameter、hk.next_rng_key）使它成为学习JAX编程模型的优秀教学工具。

📄项目信息

@software{haiku2020github,
  author = {Tom Hennigan and Trevor Cai and Tamara Norman
            and Lena Martens and Igor Babuschkin},
  title = {{H}aiku: {S}onnet for {JAX}},
  url = {http://github.com/deepmind/dm-haiku},
  version = {0.0.16},
  year = {2020},
}