Introduction

Building complex neural networks often requires a delicate balance between high-level abstractions and low-level control. For researchers and developers using JAX, managing model parameters and state across pure functions can become cumbersome. dm-haiku, a JAX-based neural network library developed by DeepMind, simplifies this process by providing an object-oriented API that integrates seamlessly with JAX’s functional transformations. With its intuitive module system, dm-haiku allows developers to define architectures without sacrificing the power of jit, grad, and vmap.

What Is dm-haiku?

dm-haiku is a neural network library built on top of JAX designed to provide simple, composable abstractions for machine learning research. It is maintained by Google DeepMind and licensed under the Apache License 2.0. The library’s primary goal is to allow users to use familiar object-oriented programming models while retaining full access to JAX’s pure function transformations.

At its core, dm-haiku provides two primary tools: the hk.Module abstraction and the hk.transform function. By wrapping a function that uses these modules, hk.transform converts a stateful-looking function into a pair of pure functions: one for initialization (init) and one for application (apply), effectively bridging the gap between object-oriented design and functional programming.

Why dm-haiku Matters

Before dm-haiku, building networks in pure JAX required manual bookkeeping of all parameters—a process that is error-prone and scales poorly for large models. dm-haiku fills this gap by automating parameter management. It preserves the programming model of Sonnet (DeepMind’s previous library for TensorFlow), which has seen near-universal adoption within DeepMind’s research projects.

For developers who prefer a minimal, modular approach to state management, dm-haiku is a powerful choice. It allows for rapid prototyping of new architectures and ensures that models are reproducible and easy to test. While the JAX ecosystem has evolved, dm-haiku remains a critical tool for those maintaining legacy DeepMind research code or those who prefer its specific approach to transform-based state management.

Key Features

• Module-Based Architecture: The hk.Module class allows users to define reusable components with their own parameters, mirroring the approach used in PyTorch or Sonnet.
• Functional Transformation: The hk.transform utility converts stateful module definitions into pure functions, ensuring compatibility with JAX’s jit and grad.
• State Management: dm-haiku provides native support for managing both parameters and mutable state (such as batch normalization statistics) through hk.transform_with_state.
• Composable Abstractions: The library offers a wide range of pre-built modules, such as hk.nets.MLP and hk.Linear, which can be composed into complex architectures.
• JAX Integration: Full compatibility with the JAX ecosystem, including libraries like Optax for optimization and Chex for testing.
• Parameter Sharing: Built-in mechanisms to share parameters across different parts of a network, which is essential for recurrent neural networks and Siamese networks.

How dm-haiku Compares

The primary alternative to dm-haiku is Flax, another JAX-based library developed by Google. While both aim to solve the same problem, they differ significantly in their programming style.

dm-haiku is often preferred by those who find the transform pattern more intuitive and closer to the “JAX spirit” of chaining functions. It is more minimal and focuses on a single, powerful abstraction. In contrast, Flax is more comprehensive and offers a wider array of features and a larger community. As of July 2023, Google DeepMind recommends that new projects adopt Flax for its more active development and extensive documentation.

Getting Started: Installation

To use dm-haiku, you must first have JAX installed. Follow the official JAX installation guide to ensure you have the correct version for your hardware (CPU, GPU, or TPU).

PyPI Installation

The simplest way to install the latest stable version of dm-haiku is via pip:

pip install -U dm-haiku

GitHub Installation

For the most recent updates and bug fixes, you can install directly from the GitHub repository:

pip install git+https://github.com/deepmind/dm-haiku

Conda Installation

If you use Conda, you can install dm-haiku from the conda-forge channel:

conda install conda-forge::dm-haiku

How to Use dm-haiku

The basic workflow in dm-haiku involves defining a function that describes your network architecture and then transforming it into a pure function using hk.transform. This allows you to separate the architecture definition from the parameter initialization and application.

First, you define a net_fn (network function) that uses Haiku modules. Then, you apply hk.transform to the this function. The resulting transformed function has two methods: init and apply. You use init to create the initial parameters based on a given input shape and an RNG key, and apply to use those parameters to perform a forward pass.

Code Examples

Below are examples of how to implement a simple Multi-Layer Perceptron (MLP) and a more complex stateful network.

Simple MLP Example

This example shows the most basic usage of dm-haiku to create a simple feed-forward network.

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

def net_fn(x):
    mlp = hk.nets.MLP([256, 128, 64], axis=-1)
    return mlp(x)

# Transform the function into a Haiku module
forward = hk.transform(net_fn)

# Initialize parameters
rng = jax.random.PRNGKey(42)
_ = forward.init(rng, jnp.ones([1, 784]))

# Apply the network
params = forward.init(rng, jnp.ones([1, 784]))
logits = forward.apply(params, None, jnp.ones([1, 784]))

Stateful Network Example

For networks that require mutable state, such as Batch Normalization, you use hk.transform_with_state. This ensures that the state is handled as a pure function, remaining compatible with JAX.

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

def net_fn(x):
    # Batch Normalization requires state
    x = hk.Linear(128)(x)
    x = hk.BatchNorm(momentum=0.9)(x, is_training=True)
    x = jax.nn.relu(x)
    return x

# Use transform_with_state for stateful modules
forward = hk.transform_with_state(net_fn)

# Initialize parameters and state
rng = jax.random.PRNGKey(42)
params, state = forward.init(rng, jnp.ones([1, 128]))

# Apply the network with state
logits = forward.apply(params, state, jnp.ones([1, 128]))

Real-World Use Cases

dm-haiku is particularly effective in scenarios where research flexibility is the primary goal. It is widely used in the following scenarios:

• Reinforcement Learning (RL): DeepMind’s RL agents often use dm-haiku to define the policy and value networks, as the modularity allows for rapid iteration on network architectures.
• Rapid Prototyping of New Architectures: Researchers can quickly test different layer types and connectivity patterns without the need to manually manage the Pytree of parameters.
• Legacy Research Codebases: Many of the most influential AI research papers from DeepMind have their official implementations in dm-haiku, making it essential for those who want to build upon those specific models.
• Custom Parameter Sharing: In Siamese networks or weight-tied architectures, dm-haiku’s module system makes it sharing weights across different calls to the same module instance easy and intuitive.

Contributing to dm-haiku

Contributions to dm-haiku are welcome, although the project is currently in maintenance mode. This means that new features are not being actively accepted, but bug fixes and compatibility updates for new JAX releases are best-effort supported.

To contribute, you can report bugs via the GitHub issue tracker. If you are submitting a PR, please ensure your changes are compatible with the existing API and focus on bug fixes or critical compatibility issues. All contributors must adhere to the the project’s code of conduct.

Community and Support

Support for dm-haiku is provided primarily through the GitHub repository. Users can find detailed documentation at the official Haiku documentation site. The community is largely composed of machine learning researchers and developers who use JAX for high-performance computing.

For those looking for a more active community and more extensive examples, the project recommends transitioning to Flax, which has a larger ecosystem of end-to-end examples and a more active development team.

Conclusion

dm-haiku is a powerful, minimal library that brings the ease of object-oriented neural network definition to the functional world of JAX. By automating parameter management through its transform pattern, it allows researchers to focus on architecture design rather than low-level bookkeeping. While it is now in maintenance mode, it remains a vital tool for those who prefer its minimal approach or are working with DeepMind’s research implementations.

If you are starting a new, large-scale project today, the recommendation is to adopt Flax. However, for those who appreciate the modularity and a minimal API, dm-haiku continues to be an excellent choice for JAX-based research. Star the repo, try the quickstart, and explore the existing research models built with it.

Resources

• dm-haiku GitHub Repository
• Official Haiku Documentation
• JAX GitHub Repository