AI & 機械学習

PyTorch デベロッパー向け JAX 基礎ガイド

2025年1月31日

Anfal Siddiqui

Machine Learning Engineer, Cloud AI

※この投稿は米国時間 2025 年 1 月 7 日に、Google Cloud blog に投稿されたものの抄訳です。

PyTorch を使われている方なら、JAX のメリットについて耳にしたことがあると思います。たとえば、その高いパフォーマンス、エレガントな関数型プログラミングアプローチ、強力な組み込みの並列計算サポートなどです。しかし、それと同時に、JAX を使い始めるときに利用できる、わかりやすく進めやすいチュートリアルを見つけるのは難しいかもしれません。JAX の新しいコンセプトを、馴染みのある PyTorch の構成要素と結び付け、JAX の基礎への理解を促すようなものです。そこでこのたび、そのようなチュートリアルを作成しました。

このチュートリアルでは、PyTorch ユーザーの視点から JAX エコシステムの基礎を探っていきます。タイタニック号の沈没事故で生き残る乗客を予測するという代表的な ML タスクを取り上げ、簡単なニューラルネットワークを両方のフレームワークでトレーニングします。その過程で、モデルの定義やインスタンス化、トレーニングなど、JAX の多くの要素と、PyTorch の同等の要素の対応を示しながら、JAX を紹介していきます。

付属のノートブックで、コード例の全文をご確認いただけます（https://www.kaggle.com/code/anfalatgoogle/pytorch-developer-s-guide-to-jax-fundamentals）。

JAX におけるモジュール性

PyTorch ユーザーから見て、高度にモジュール化された JAX のエコシステムは、使い慣れた PyTorch とはかなり違うと最初は感じるかもしれません。JAX は、自動微分をサポートする、高パフォーマンスの数値計算ライブラリを目指しています。PyTorch とは異なり、ニューラルネットワーク、オプティマイザなどの定義のサポートを明示的に組み込もうとはせず、柔軟性を念頭に設計されているため、任意のフレームワークを機能に追加できるようになっています。

このチュートリアルでは、Flax ニューラルネットワークライブラリと Optax 最適化ライブラリを使用します。いずれも一般的に使われ、サポートも充実したライブラリです。先に、新しい Flax NNX API でニューラルネットワークをトレーニングする方法と、PyTorch との類似性を確認します。その後、同じことを、現在も広く使われている以前の Linen API で行う方法を示します。

関数型プログラミング

チュートリアルを始める前に、PyTorch やその他のフレームワークで使われているオブジェクト指向プログラミングではなく、関数型プログラミングが JAX で使われている理由を説明します。簡単に言うと、関数型プログラミングでは、状態を変更できず、副作用を生じさせない、つまり同じ入力に対して常に同じ出力を生成する、純粋関数を主に使います。JAX では、コンポーズ可能な関数と変更できない配列の多用という形でこれが現れています。

純粋関数と関数型プログラミングの予測しやすいという性質が、JAX の多くのメリットを可能にしています。たとえば、ジャストインタイム（JIT）コンパイルにより、XLA コンパイラは GPU や TPU のコードを大幅に最適化し、速度を大きく向上できます。また、JAX では処理のシャーディングや並列化が非常に簡単になっています。詳細については、JAX の公式チュートリアルをご覧ください。

関数型プログラミングに慣れていなくても心配ご無用です。これから見ていくように、Flax NNX ではこうしたことが標準の Python に似たイディオムの背後に隠されています。

データ読み込み

JAX におけるデータ読み込みは非常にシンプルで、PyTorch と同じように行います。PyTorch のデータセット / データローダを使用して、簡単な collate_fn で、JAX のすべての計算を支える NumPy に似た配列への変換を行うことができます。

読み込んでいます...

import torch
from torch.utils.data import Dataset, DataLoader, default_collate
import jax.numpy as jnp
from jax.tree_util import tree_map

# データセット定義
class TitanicDataset(Dataset):
    def __init__(self, samples, labels):
        self.df = samples
        self.labels = labels

def __len__(self):
        return len(self.df)

def __getitem__(self, idx):
        x = torch.tensor(self.df.iloc[idx].values, dtype=torch.float32)
        y = torch.tensor(self.labels.iloc[idx], dtype=torch.float32)
        return x, y

def numpy_collate(batch):
  return tree_map(jnp.asarray, default_collate(batch))

# pandas データフレームから Train データセット、データローダを作成
train_dataset = TitanicDataset(X_train, y_train)
train_dataloader_jax = DataLoader(train_dataset, batch_size=64, shuffle=True, collate_fn=numpy_collate)
train_dataloader_torch = DataLoader(train_dataset, batch_size=64, shuffle=True)

# pandas データフレームから Eval データセット、データローダを作成
eval_dataset = TitanicDataset(X_test, y_test)
eval_dataloader_jax = DataLoader(eval_dataset, batch_size=64, shuffle=False, collate_fn=numpy_collate)
eval_dataloader_torch = DataLoader(eval_dataset, batch_size=64, shuffle=False)

モデルの定義

Flax NNX API におけるニューラルネットワークの定義は、PyTorch の場合とよく似ています。ここでは、簡単な 2 層の多層パーセプトロンを両方のフレームワークで定義します。まずは PyTorch からです。

読み込んでいます...

NNX のモデル定義は上記の PyTorch コードとよく似ています。いずれの場合も __init__ を使用してモデルの層を定義します。__call__ は forward に対応します。

読み込んでいます...

モデルの初期化と使用

NNX でのモデルの初期化は PyTorch とほぼ同じです。いずれのフレームワークでも、モデルクラスのインスタンスをインスタンス化するときに、モデルのパラメータが（遅延初期化ではなく）早期初期化され、インスタンス自体と関連付けられます。NNX の唯一の違いは、モデルのインスタンス化のときに疑似乱数生成器（PRNG）キーを渡す必要があることです。NNX では、JAX の関数型という性質に沿って、暗黙的なグローバルランダム状態が回避されるため、PRNG キーを明示的に渡す必要があります。これにより、PRNG の生成が容易に再現可能、並列化可能、ベクトル化可能になります。詳細については、JAX のドキュメントをご覧ください。

読み込んでいます...

モデルを実際に使ったデータ群の処理は、2 つのフレームワークで同じです。

読み込んでいます...

トレーニングのステップと誤差逆伝播

トレーニングループに関しては、PyTorch と Flax NNX とで、いくつか重要な違いがあります。それを確認するために、NNX の完全なトレーニングループを順を追って組み立ててみましょう。

設定

読み込んでいます...

いずれのフレームワークでも、オプティマイザを作成し、独自の最適化アルゴリズムを柔軟に指定できます。PyTorch ではモデルのパラメータを渡す必要がありますが、Flax NNX ではモデルを直接渡すだけで、背後にある Optax オプティマイザですべてのインタラクションが処理されます。

フォワード + バックワード パス

読み込んでいます...

PyTorch と JAX の最大の違いは、おそらく完全なフォワード / バックワードパスを行う方法です。PyTorch では loss.backward() で勾配を計算します。このとき AutoGrad がトリガーされ、loss の計算グラフに沿って勾配が計算されます。

これに対して、JAX の自動微分はより数学に近く、関数の勾配があります。具体的には、nnx.value_and_grad / nnx.grad で関数 loss_fn を受け取り、関数 grad_fn を返します。その後、grad_fn 自体が、入力に応じて、loss_fn の出力の勾配を返します。

この例では、PyTorch と同じことを loss_fn で行っています。最初にフォワードパスから logits を取得してから、おなじみの loss を計算しています。その後 grad_fn が、model のパラメータに応じて、loss の勾配を計算しています。数学的には、返される grads は ∂J/∂θ です。これは PyTorch の背後で行われている処理と同じです。PyTorch では loss.backward() の実行時に勾配がテンソルの .grad 属性に格納されるのに対して、JAX / Flax NNX では、状態を変更しないという関数型アプローチに従い、勾配が直接返されています。

オプティマイザのステップ

読み込んでいます...

PyTorch では、optimizer.step() により、勾配を使用して重みがインプレースで更新されます。NNX でも重みがインプレースで更新されますが、バックワードパスで計算した grads を直接渡す必要があります。これは PyTorch と同じ最適化のステップですが、JAX の基礎にある関数型という性質に沿って、PyTorch よりもわずかに明示的になっています。

完全なトレーニング ループ

これで、JAX / Flax NNX で完全なトレーニングループを構築するために必要なものがすべて揃いました。参考として、よく知っている PyTorch のループを先に確認しましょう。

読み込んでいます...

次に、NNX の完全なトレーニングループです。

読み込んでいます...

import optax

def train(model, train_dataloader, eval_dataloader, num_epochs):
  optimizer = nnx.Optimizer(model, optax.adam(learning_rate=0.01))

for epoch in (pbar := tqdm(range(num_epochs))):
    pbar.set_description(f"Epoch {epoch}")
    model.train()
    for batch in train_dataloader:
      train_step(model, optimizer, batch)

pbar.set_postfix(train_accuracy=eval(model, train_dataloader),    eval_accuracy=eval(model, eval_dataloader))

@nnx.jit
def train_step(model, optimizer, batch):
  def loss_fn(model):
    logits = model(batch[0])
    loss = optax.sigmoid_binary_cross_entropy(logits.squeeze(),batch[1]).mean()
    return loss
  grad_fn = nnx.value_and_grad(loss_fn)
  loss, grads = grad_fn(model)
  optimizer.update(grads)

def eval(model, eval_dataloader):
  model.eval()
  total = 0
  num_correct = 0
  for batch in eval_dataloader:
    res = eval_step(model, batch)
    total += res.shape[0]
    num_correct += jnp.sum(res)
  return num_correct / total

@nnx.jit
def eval_step(model, batch):
  logits = model(batch[0])
  logits = logits.squeeze()
  preds = jnp.round(nnx.sigmoid(logits))
  return preds == batch[1]

最大のポイントは、トレーニングループが PyTorch と JAX / Flax NNX とでよく似ているということです。相違点のほとんどは、オブジェクト指向プログラミングか関数型プログラミングかの違いに行き着きます。関数型プログラミングと、関数の勾配を考えるという点については、慣れるまでに少し時間がかかるかもしれませんが、慣れてしまえば JIT コンパイルや自動並列化といった前述の JAX のメリットを享受できます。たとえば、@nnx.jit のアノテーションを上記の関数に追加するだけで、P100 GPU で 500 エポックのモデルのトレーニング時間を 6.25 分からわずか 1.8 分に短縮できることが Kaggle で確認されています。他の CPU や TPU、NVIDIA 以外の GPU でも同程度のスピードアップを見込めることがわかっています。

Flax Linen のリファレンス

前述のとおり、JAX のエコシステムは非常に柔軟であり、任意のフレームワークを使用できるようになっています。新規のユーザーには NNX がおすすめのソリューションですが、Flax Linen API も、MaxText や MaxDiffusion などの強力なフレームワークで現在も広く使用されています。NNX はより Python に似て、状態管理の複雑さが隠されていますが、Linen は純粋関数型プログラミングにより忠実になっています。

JAX のエコシステムに加わる場合、この両方に慣れておくことが大きなメリットになります。その一助となるよう、以下では NNX のほとんどのコードを Linen で再現し、主な違いをコメントとして入れます。

読み込んでいます...

# 完全なトレーニング ループ 
import optax
from flax.training import train_state

# TrainState はパラメータと勾配の管理に役立つコンビニエンス ラッパー  
optimizer = optax.adam(learning_rate=0.01)

state = train_state.TrainState.create(
    apply_fn=flax_model.apply,
    params=params,
    tx=optimizer,
)

def train(state, train_dataloader, eval_dataloader, subdropout, num_epochs):
  for epoch in (pbar := tqdm(range(num_epochs))):
    pbar.set_description(f"Epoch {epoch}")
    for batch in train_dataloader:
      state, loss = train_step(state, batch, subdropout)

pbar.set_postfix(train_accuracy=eval(state, train_dataloader), eval_accuracy=eval(state, eval_dataloader))

return state

def eval(state, eval_dataloader):
  total = 0
  num_correct = 0
  for batch in eval_dataloader:
    res = eval_step(state, batch)
    total += res.shape[0]
    num_correct += jnp.sum(res)
  return num_correct / total

@jit
def train_step(state, batch, subdropout):
  def loss_fn(params):
    logits = state.apply_fn(params, batch[0], True, rngs={'dropout': subdropout})
    loss = optax.sigmoid_binary_cross_entropy(logits.squeeze(), batch[1]).mean()
    return loss

grad_fn = jax.value_and_grad(loss_fn)
  loss, grads = grad_fn(state.params)
   # grads を TrainState に渡し、パラメータが更新された新しい TrainState を関数型プログラミング スタイルで取得
  state = state.apply_gradients(grads=grads)
  return state, loss

@jit
def eval_step(state, batch):
  logits = state.apply_fn(state.params, batch[0], False)
  logits = logits.squeeze()
  preds = jnp.round(nn.sigmoid(logits))
  return preds == batch[1]

次のステップ

このブログ投稿で得た JAX / Flax の知識により、独自のニューラルネットワークを記述する準備ができたと思います。Google Colab か Kaggle を使用すれば、すぐに始められます。Kaggle で課題を探し、Flax NNX を使ってモデルを新規に記述するか、MaxText を使って大規模言語モデル（LLM）のトレーニングを開始するなど、可能性は無限にあります。

今回取り上げたのは、JAX と Flax の一部にすぎません。JIT、自動ベクトル化、カスタム勾配などの詳細については、JAX と Flax のドキュメントをご覧ください。

-Cloud AI、ML エンジニア Anfal Siddiqui

投稿先

AI & Machine Learning

Vertex AI Gen AI Evaluation Service でのエージェント評価のご紹介

執筆者: Irina Sigler • 所要時間: 7 分

https://storage.googleapis.com/gweb-cloudblog-publish/images/image_bolg_2436_x1200.max-700x700.png

AI & Machine Learning

[3月 13日（木）初開催] AI Agent Summit ’25 Spring

執筆者: Google Cloud Japan Team • 所要時間: 2 分

Compute

新年、AI Hypercomputer を新たにアップデート

執筆者: Chelsie Czop • 所要時間: 10 分

https://storage.googleapis.com/gweb-cloudblog-publish/images/hero_image_jalinfotec_horizontal.max-700x700.jpg

Customers

JALインフォテック: オンプレミス刷新と生成 AI 導入による Lift & Transform で、業務効率化や人材育成などを推進

執筆者: Google Cloud Japan Team • 所要時間: 4 分

PyTorch デベロッパー向け JAX 基礎ガイド

Anfal Siddiqui

JAX におけるモジュール性

関数型プログラミング

データ読み込み

モデルの定義

モデルの初期化と使用

トレーニングのステップと誤差逆伝播

Flax Linen のリファレンス

次のステップ

関連記事

Vertex AI Gen AI Evaluation Service でのエージェント評価のご紹介

[3月 13日（木）初開催] AI Agent Summit ’25 Spring

新年、AI Hypercomputer を新たにアップデート

JALインフォテック: オンプレミス刷新と生成 AI 導入による Lift & Transform で、業務効率化や人材育成などを推進