今週、ZyphraはApache 2.0ライセンスの下でZUNA1.1をリリースしました。この脳波(EEG)基盤モデルは、任意のチャンネルレイアウトでデータを再構成、ノイズ除去、アップサンプリングします。これはZyphraの以前のオープンEEG基盤モデルZUNA1をベースにしています。
主な変更点は、精度の向上ではなく柔軟性です。実際のEEG記録は乱雑です。セッションの長さは異なり、チャンネルは途中でノイズが入ったりドロップしたりします。モンタージュは4電極ヘッドバンドから256チャンネル研究用キャップまでさまざまです。ZUNA1は固定5秒のセグメントのみ処理できましたが、ZUNA1.1は0.5~30秒の可変長入力を受け付けます。
ZUNA1.1とは
この柔軟性を理解するには、モデルの機能から始めましょう。ZUNA1.1は、頭皮EEG信号向けの3億8000万パラメータのマスク拡散オートエンコーダです。チャンネルのサブセットが与えられると、既存のEEGセグメントとチャンネルをノイズ除去し、欠落したチャンネルを再構成します。また、頭皮上の物理座標に基づいて新しいチャンネル信号を予測することもできます。
パラメータ数はZUNA1から変更なく、コンシューマーGPUで動作し、多くのワークロードでCPUでも十分に機能します。重みはHugging Faceに、推論および前処理コードはGitHubにあります。pip install zunaでインストール可能です。Zyphraは無料のブラウザベースEEGプレイグラウンドもホストしており、これらはすべて研究目的に限り提供されています。
アーキテクチャの仕組み
この柔軟性はトークン化に基づいています。ZUNAはトランスフォーマーエンコーダ-デコーダ拡散オートエンコーダです。各チャンネルを0.125秒のセグメント(256Hzで32サンプル)にスライスし、各セグメントは連続値トークンになります。トークンはチャンネル×時間の順にシリアル化されます。
位置符号化が鍵です。各トークンは(x, y, z, t)の4D回転位置符号化を持ちます。これは電極の3D頭皮座標と粗い時間インデックスです。配列インデックスではなく位置がモデルにチャンネルの場所を伝えるため、ZUNAはチャンネルに依存しません。任意の電極レイアウトを受け入れ、記録されたことのない位置の信号も生成できます。この能力により、位置に基づく任意のチャンネルアップサンプリングが可能になります。
エンコーダは信号を潜在表現に圧縮し、その潜在表現が適応RMSノルムを介してデコーダを条件付けます。デコーダは整流流目的関数で訓練されます。ZUNA1.1のアーキテクチャ変更は主に訓練の安定性を目的としており、正規化層の追加などが含まれます。
ZUNA1からの変更点
アーキテクチャは近いため、違いは訓練に起因します。
- 可変長入力(0.5~30秒):ZUNA1.1は訓練例ごとにセグメント長をサンプリングし、0.125秒のトークングリッドにスナップします。長さは4つのビンから抽出され、非常に短いものから長いものまでカバーします。中間の1.5~10秒範囲は最も一般的な動作点であるためオーバーサンプリングされます。トークン数が変動するため、Zyphraは固定予算まで複数のセグメントを1バッチにパックします。サンプル認識マスクを備えたフレキシブルアテンションにより、トークンがサンプル間で相互作用するのを防ぎます。そのため、1つのモデルで0.5秒のスニペットと30秒の長尺を再構成なしに処理できます。
- より豊富な再構成タスクの混合:ZUNA1は1つのドロップアウトパターン(一様ランダムな全チャンネル)のみで訓練されました。ZUNA1.1は4つのパターンで訓練されます。1つ目は全チャンネルドロップアウトで、スパースモンタージュや死んだ電極をカバーします。2つ目はすべてのチャンネルから短い時間範囲を削除します。3つ目は一部のチャンネルからのみそれらの範囲を削除し、空間と時間にギャップを集中させます。4つ目は個々のポイントに欠損値を散在させます。
- 品質認識前処理と大規模コーパス:ZUNA1はチャンネル品質の判断を記録全体レベルで行い、使用可能な信号を破棄していました。ZUNA1.1は代わりにチャンネルごと、秒ごとの品質スコアを計算し、読み込み時に閾値処理します。これにより、コーパスは公開EEGデータの約200万チャンネル時間から約350万チャンネル時間に拡大しました。Zyphraチームはまた、記録ごとに2つのフィルタバリアント(0.1–45 Hzバンドパス、および0.01 Hzハイパス+ノッチ)を事前計算します。前処理戦略の一般化は目標として掲げられていますが、まだベンチマーク結果はありません。
結果
したがって、柔軟性が精度を犠牲にしたかが問題です。保持タスクでは、ZUNA1.1はZUNA1と同等またはより良い再構成NMSEを達成しました。両者ともMNEの古典的な球面スプライン補間を明確に上回っています。公平な比較のため、これらの評価セットは正確に5秒のサンプルを使用しました。
Zyphraは領域ベースのテストも実施しました。1つの脳領域の電極を削除し、残りの7領域から再構成します。この設定はランダムなチャンネルドロップよりも現実的です。ZUNA1.1は球面スプラインとZUNA1の両方を上回りました。
インタラクティブな説明
これらのメカニズムを具体的にするため、以下のデモでエンドツーエンドのパイプラインをアニメーション化しています。
ZUNA1 vs ZUNA1.1
総合すると、2つのリリースの違いはアーキテクチャではなく訓練にあります。
| 属性 | ZUNA1 | ZUNA1.1 | |------|-------|---------| | パラメータ | 380M | 380M | | アーキテクチャ | トランスフォーマーエンコーダ-デコーダ拡散オートエンコーダ | 同一、追加の正規化層 | | 入力長 | 固定5秒 | 0.5~30秒、0.125秒グリッドにスナップ | | トークン | 0.125秒/32サンプル(256Hz) | 同一 | | 位置符号化 | 4D RoPE (x,y,z,t) | 同一 | | デコーダ目的関数 | 整流流 | 整流流 | | 訓練中のドロップアウトスキーム | 1種類(一様ランダム全チャンネル) | 4種類(チャンネル、時間、チャンネル×時間、散在) | | 訓練コーパス | ~200万チャンネル時間 | ~350万チャンネル時間 | | 品質フィルタリング | 全記録レベル | チャンネルごと秒ごとのスコア、読み込み時閾値 | | 前処理バリアント | 単一 | 2種類(0.1–45 Hzバンドパス;0.01 Hzハイパス+ノッチ) | | ライセンス | Apache 2.0 | Apache 2.0 | | 再構成NMSE | ベースライン | 同等以上 |
実行方法
実践では、reconstruct_fifは.fifファイルを直接処理し、.pt形式を経由しません。従来の4ステップパイプラインも引き続き提供されています。
from zuna import reconstruct_fif
reconstruct_fif(
input_dir="fif_in",
output_dir="fif_out",
figures_dir="figures",
gpu_device=0, # GPU ID、または空文字でCPU
segment_sec=5.0, # ウィンドウ長;デフォルトは5.0、完全な30秒ではない
montage="standard_1020", # フォールバック、ファイルに位置がない場合のみ使用
repair_channels=["Cz"], # 完全再構成するチャンネル
target_channel_count=["Fz", "Pz"], # 名前で新しいチャンネルを追加/アップサンプリング(intで自動選択も可)
bad_segments=[(5, 6), (10, 11, "C3")], # 不良時間範囲をマーク(全チャンネル、または1チャンネル)
sample_steps=50, # 拡散ステップ数
)デフォルトに注意: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")]を渡し、1チャンネルの1期間をクリーン。 - ヘッドバンドアップサンプリング:4つの電極を入力し、追加で標準の
standard_1005位置を要求。 - UI駆動のクリーニング:
mask_dirでファイルごとにマスクを提供し、ターゲットに結合。
モデル重み、GitHubリポジトリ、技術詳細をチェックしてください。また、Twitterでフォローしていただき、15万以上のML SubRedditに参加し、ニュースレターを購読することをお忘れなく。Telegramをご利用ですか?Telegramでも参加いただけます。
GitHubリポジトリやHugging Faceページ、製品リリース、ウェビナーなどのプロモーションで協力が必要ですか?お問い合わせください。
この記事「Zyphra、可変長入力をサポートするApache 2.0 EEG基盤モデルZUNA1.1をリリース」はMarkTechPostに最初に掲載されました。