Penzai: 谷歌DeepMind推出的JAX神经网络可视化与编辑工具包

Penzai:让神经网络变得更加透明可控

在人工智能飞速发展的今天,深度神经网络已经成为许多领域不可或缺的技术工具。然而,这些复杂的模型往往被视为"黑盒",其内部运作机制难以理解和解释。为了让研究人员能够更好地探究和操控神经网络模型,谷歌DeepMind最近推出了一个名为Penzai的开源工具包。

Penzai的由来与设计理念

Penzai这个名字源自中国古老的盆栽艺术,寓意着将复杂的事物缩小并精心塑造。正如盆栽艺术家将自然景观浓缩在小小的盆中一样,Penzai的设计目标是将庞大复杂的神经网络模型变得易于理解和操作。

这个工具包基于JAX构建,专注于让研究人员能够在模型训练完成后对其进行各种操作。无论是进行逆向工程、消融实验、内部激活探测,还是模型手术和架构调试,Penzai都提供了强大而灵活的工具支持。

Penzai的主要功能模块

Penzai由一系列模块化工具组成,每个模块都可以独立使用,也可以协同工作:

Treescope: 这是一个功能强大的交互式Python美化打印工具,可以直观地展示深度嵌套的JAX pytree结构,内置了对任意维度NDArray的可视化支持。
penzai.core.selectors: 这是一个pytree瑞士军刀,将JAX的.at[...].set(...)语法推广到任意类型驱动的pytree遍历中,使得对Penzai模型和其他数据结构进行复杂重写或即时修补变得容易。
penzai.core.named_axes: 这是一个轻量级的命名轴系统,可以将普通的JAX函数提升为在命名轴上进行向量化操作,让用户能够在命名和位置编程风格之间无缝切换。
penzai.nn: 这是一个声明式的、基于组合子的神经网络库,将模型表示为易于修改的数据结构。它暴露了模型前向传播的完整结构,并允许在树的叶节点包含可变变量,以跟踪可变状态和参数共享。
penzai.models.transformer: 这是一个模块化实现的常见Transformer架构,支持加载Gemma、Llama、Mistral和GPT-NeoX/Pythia等预训练权重。它使用模块化组件和命名轴构建,简化了复杂的模型操作工作流程。

Penzai的使用示例

让我们通过一个简单的例子来看看Penzai是如何工作的:

import penzai as pz
import jax

# 初始化一个简单的神经网络
from penzai.models import simple_mlp
mlp = simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(0),
    feature_sizes=[8, 32, 32, 8]
)

# 模型结构会自动可视化
mlp

# 捕获和提取元素级非线性后的激活
@pz.pytree_dataclass
class AppendIntermediate(pz.nn.Layer):
  saved: pz.StateVariable[list[Any]]
  def __call__(self, x: Any, **unused_side_inputs) -> Any:
    self.saved.value = self.saved.value + [x]
    return x

var = pz.StateVariable(value=[], label="my_intermediates")

# 创建一个保存激活的模型副本
saving_model = (
    pz.select(mlp)
    .at_instances_of(pz.nn.Elementwise)
    .insert_after(AppendIntermediate(var))
)

output = saving_model(pz.nx.ones({"features": 8}))
intermediates = var.value