DeepMind Haiku如何追踪网络层、维持参数一致性

核心底层逻辑 #

你观察到的「每次前向调用重新实例化网络层但参数保持一致」，是Haiku作为函数式神经网络框架的核心设计特性：

• 你代码里实例化的hk.Embed、hk.Linear、model.Transformer对象本身不存储任何参数，它们只是携带层配置（比如输出维度、初始化方式）的轻量临时对象，生命周期仅在单次前向函数调用内。
• Haiku会在被hk.transform包裹的前向函数运行时，维护一个隐式的全局上下文：
  1. 首次运行初始化（init）流程时，Haiku会沿着前向调用的栈结构，给每个模块、每个直接通过hk.get_parameter创建的参数生成全局唯一的层级路径名，同时调用你指定的初始化函数生成参数张量，统一存入上下文绑定的参数字典中，路径作为键、参数作为值。
  2. 后续每次前向推理、梯度计算调用前向函数时，只要层的实例化顺序、嵌套层级和初始化时完全一致，Haiku就会生成和初始化阶段完全相同的路径名，直接从参数字典里取已经存在的参数绑定到当前临时创建的层对象上，不会重新初始化参数。
     你代码里的位置嵌入直接通过hk.get_parameter('pos_embs', ...)创建，本质也是走的同样的路径匹配逻辑，和层对象创建的参数没有区别。

获取权重、自定义操作的实现方式 #

你不需要获取层对象的可变引用——因为层对象本身根本不持有参数，所有可读写的参数都存在hk.transform返回的init函数输出的参数字典中，这是个普通的嵌套JAX数组字典，可以直接遍历、打印、修改：

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

# 构造前向函数并做函数式转换
forward_fn = build_forward_fn(
    vocab_size=32000, d_model=512, num_heads=8,
    num_layers=6, dropout_rate=0.1
)
transformed_forward = hk.transform(forward_fn)

# 初始化参数（仅需运行一次）
rng_key = jax.random.PRNGKey(0)
dummy_input = {"obs": jnp.ones((2, 128), dtype=jnp.int32)}  # batch=2, seq_len=128
params = transformed_forward.init(rng_key, dummy_input, is_training=True)

# 直接操作参数字典即可完成权重打印、修改等操作
# 打印token嵌入层权重
print("Token embedding matrix shape:", params["embed"]["w"].shape)
# 打印位置编码参数
print("Positional embedding shape:", params["pos_embs"].shape)
# 打印最终投影层权重
print("Output projection weight shape:", params["linear"]["w"].shape)

# 可以遍历整个参数字典打印所有层的参数信息
for module_path, module_params in params.items():
    print(f"\nModule path: {module_path}")
    for param_name, param_value in module_params.items():
        print(f"  Param {param_name}, shape: {param_value.shape}, mean: {jnp.mean(param_value):.4f}")

注意：不要尝试缓存单次前向里创建的hk.Embed、hk.Linear等层对象跨前向调用复用，Haiku的上下文在每次前向调用结束后会重置，缓存层对象极易触发模块名冲突、参数匹配错误的问题，官方示例中每次前向重新实例化层的写法是标准正确用法。
如果需要在前向执行过程中获取中间激活、拦截层计算逻辑，可以使用Haiku内置的hk.intercept_methods或hk.Module的hook方法实现，不需要持有层对象引用。

内容的提问来源于stack exchange，提问作者Foobar