Zyphra が EEG 基盤モデル「ZUNA1.1」を Apache 2.0 で公開

今週、Zyphra は Apache 2.0 ライセンスの下で「ZUNA1.1」を公開しました。この EEG（脳波）基盤モデルは、任意のチャネル構成に対応し、データの再構築、ノイズ除去、アップサンプリングを実現します。これは Zyphra が以前にオープンソース化した ZUNA1 をベースに開発されたものです。

今回の主な進化は、精度の飛躍的な向上ではなく「柔軟性」にあります。実際の脳波記録は常にノイズが多く、セッションの長さは一定ではありません。また、測定中にチャネルがノイズを帯びたり切断されたりすることも珍しくありません。使用する電極の数も、4 電極のヘッドバンドから 256 チャンネルの研究用キャップまで様々です。従来の ZUNA1 は固定された 5 秒間のセグメントしか処理できませんでしたが、ZUNA1.1 では 0.5 秒から 30 秒までの可変長入力をサポートしています。

ZUNA1.1 とは何か？

その柔軟性を理解するには、まずこのモデルが何をするものかを知る必要があります。

ZUNA1.1 は、スカルプ EEG（頭皮脳波）信号を対象とした、380M パラメータのマスク付き拡散オートエンコーダーです。入力された一部のチャネルデータから既存の脳波セグメントやチャネルのノイズを除去し、欠落しているデータを復元します。さらに、頭上の物理的な座標情報に基づいて、存在しない新しいチャネル信号も予測することが可能です。

パラメータ数は ZUNA1 と同じです。コンシューマー向けの GPU で動作可能であり、多くのワークロードでは CPU でも十分な性能を発揮します。モデルの重みは Hugging Face に、推論および前処理コードは GitHub に公開されています。インストールは pip install zuna で完了します。また、Zyphra はブラウザ上で利用できる無料の EEG プレイグラウンドも提供しており、これらはすべて研究利用を目的として無償で配布されています。

アーキテクチャの仕組み

この柔軟性の根幹にあるのが「トークナイゼーション」です。

ZUNA は、トランスフォーマーエンコーダー・デコーダーを基盤とした拡散オートエンコーダーです。各チャネルの信号は 0.125 秒（サンプリング周波数 256Hz の場合 32 サンプル）ごとにスライスされ、連続値を持つトークンとして扱われます。これらのトークンは、チャネル順に時間軸に沿って並べられます。

ここで重要なのが位置エンコーディングです。各トークンは、電極の 3D スカルプ座標（x, y, z）と粗い時間インデックス（t）を含む 4 次元ロータリー位置エンコーディングを保持します。モデルがチャネルの位置を把握するのは配列のインデックスではなく、物理的な位置情報だからです。この仕組みにより ZUNA はチャネル構成に依存せず、任意の電極レイアウトを受け入れることができます。さらに、記録されたことのない位置でも信号を生成することが可能で、これによって場所に基づく任意のチャネルアップサンプリングが可能になります。

エンコーダーは信号を潜在表現（latent）に圧縮し、その潜在変数は適応型 RMS ノルムを通じてデコーダーを制御します。デコーダーは整流フロー（rectified-flow）の目的関数を用いて訓練されます。ZUNA1.1 では、追加された正規化層などにより、学習の安定性を高めるためのアーキテクチャ上の改良が施されました。

ZUNA1 からの変更点

アーキテクチャ自体は大きく変わっていません。主な違いは訓練プロセスにあります。

1. 可変長入力（0.5〜30 秒）：ZUNA1.1 は、トレーニング例ごとにサンプリングされたセグメント長を 0.125 秒ごとのトークングリッドにスナップします。この長さは非常に短いものから長いものまで 4 つのビンに分類され、その中からランダムに選ばれます。特に最も一般的な運用ポイントとなる 1.5〜10 秒の範囲は過剰サンプリングされています。トークン数が可変であるため、Zyphra は固定されたバジェット上限まで、1 つのバッチ内に複数のセグメントをパッキングします。サンプル認識型のマスクを用いた Flex Attention（フレックスアテンション）により、異なるサンプル間のトークンが相互に注意を向けるのを防ぎます。これにより、再設定を行うことなく、0.5 秒のスニペットから 30 秒のストッチまで、単一のモデルで幅広く対応することが可能になります。

1. より多様な再構成タスクのミックス：ZUNA1 は一様なランダムな全チャンネルドロップアウトという 1 つのパターンでのみ学習していましたが、ZUNA1.1 では 4 つの異なるパターンを組み合わせます。最初の手法は、スパースなマウンテージや死んだ電極をカバーするための「全チャンネルドロップアウト」です。2 つ目は、すべてのチャンネルで短い時間区間を削除する手法です。3 つ目は、特定のチャンネルのみからその区間を削除し、空間的・時間的にギャップをクラスタリングさせるものです。そして 4 つ目は、欠損値を個々のポイントに散らす手法です。

1. 品質を考慮した前処理と、より大規模なコーパス

ZUNA1 は録音全体レベルでチャネルの品質判定を行っていたため、利用可能な信号まで捨ててしまうという課題がありました。一方、ZUNA1.1 ではチャネルごとに、かつ秒単位で品質スコアを算出し、ロード時に閾値を設定します。これにより、公開されている EEG データの corpus は約 200 万チャネル時間から約 350 万チャネル時間に拡大されました。

Zyphra チームはさらに、録音ごとに 2 つのフィルタバリアントを事前計算しています。1 つは 0.1〜45 Hz のバンドパスフィルタ、もう 1 つは 0.01 Hz のハイパスフィルタにノッチフィルタを組み合わせたものです。

前処理戦略間での一般化が可能であることは、ベンチマーク結果として示された事実というより、明確な目標として掲げられています。

結果

つまり、柔軟性を追求したことで精度が犠牲になったのかどうかが問われます。

検証用タスク（ホールドアウト）において、ZUNA1.1 は ZUNA1 と同等かそれ以上の再構成 NMSE を達成しました。両モデルとも、MNE で使われる古典的な球面スプライン補間法を明確に上回っています。公平な比較のために、これらの評価セットではすべて 5 秒間のサンプルを使用しています。

Zyphra はまた、領域ベースのテストも実施しました。これはある脳領域からの電極データを削除し、残りの 7 つの電極データから再構成するものです。この設定は、ランダムにチャネルを欠落させる手法よりも現実的です。その結果、ZUNA1.1 は球面スプライン法および ZUNA1 を上回る性能を示しました。

インタラクティブな解説

これらの仕組みを具体的に理解してもらうため、以下のデモではパイプラインの全体像をアニメーションで示しています。

ZUNA1 と ZUNA1.1 の比較

これらのリリースの主な違いは、アーキテクチャではなく学習プロセスにあります。

実際に動かすには、reconstruct_fif を使用します。これは .pt ファイルを経由する従来の手順を踏まずに、直接 .fif ファイルを処理できます。なお、従来の 4 ステップパイプラインも併用可能です。

from zuna import reconstruct_fif

※異なるブラウザを使用してください

reconstruct_fif(

input_dir="fif_in",

output_dir="fif_out",

figures_dir="figures",

gpu_device=0, # GPU id, or "" for CPU

segment_sec=5.0, # window length; default is 5.0, not the full 30 s

montage="standard_1020", # fallback, used only if the file has no positions

repair_channels=["Cz"], # channel(s) to fully reconstruct

target_channel_count=["Fz", "Pz"], # add/upsample new channels by name (or an int for auto)

bad_segments=[(5, 6), (10, 11, "C3")], # mark time spans bad (all channels, or one)

sample_steps=50, # diffusion steps; note: not "diffusion_sample_steps"

)

デフォルト値に注意してください。segment_sec は 5.0 に設定されているため、0.5〜30 秒の範囲を扱うには明示的な指定が必要です。電極位置はファイル自体から読み込まれます。montage 引数は、ファイル内に位置情報がない場合のみフォールバックとして使用され、3D 座標を持たないチャネルは除外されます。

再構成の対象となるのは、結合された集合です。これは、ファイル内の MNE が定義する不良チャネルや BAD_ アノテーションに、上記で指定した項目をすべて加えたものになります。出力先には 2 つのディレクトリが作成されます。full_reconstruction/ にはモデルによる再構成結果が全域に保存され、hybrid/ には元のデータと推測されたセルのみを補完したデータ、および _mask.npz が格納されます。

使用例と具体案

マスキング機能が柔軟になったことで、いくつかの実用的なパターンが利用可能になりました。

死んだ電極の処理：repair_channels=["Cz"] を指定して、周囲の電極からそのチャンネルを再構築します。

試行中のモーションアーティファクトの除去：bad_segments=[(10, 11, "C3")] を渡すことで、特定のチャンネル上の特定の区間だけをクリーンアップできます。

ヘッドバンドのアップサンプリング：4 つの電極を入力として与え、標準的な 10-05 配置の追加位置をリクエストします。

UI ベースのクリーニング：mask_dir を通じてファイルごとのマスク情報を提供し、ターゲットに統合します。