TrackCore-F: Deploying Transformer-Based Subatomic Particle Tracking on FPGAs

本論文は、FPGA 上で Transformer ベースのサブ原子粒子追跡モデルを効率的に推論可能にするための、単一または分割合成の手法とツールの開発、および TrackFormers プロジェクトに由来するモデルを用いた予備結果の提示について述べています。

要約

🌌 物語の舞台：巨大な粒子の迷路

まず、背景から説明しましょう。
スイスにある**LHC（大型ハドロン衝突型加速器）**という、地球の裏側を一周するほどの巨大なトンネルがあります。ここでは、粒子を光速に近い速さでぶつけ合い、宇宙の誕生の瞬間を再現しています。

この衝突で無数の「粒子」が飛び散ります。これを**「粒子の軌跡（トラック）」と呼びますが、まるで「満員電車の中で、数百人の乗客が同時に動き回り、誰が誰と並走しているかを探す」**ような難易度です。

これまでの研究では、この「誰が誰と並走しているか」を見つける作業は、**「後でゆっくり分析する（ポストモーテル）」**ものでした。計算量が膨大すぎて、その場ですぐに答えを出すのが難しかったからです。

🚀 目指すもの：FPGA 上の「超高速 AI 探偵」

そこで登場するのが、この論文の主人公たちです。

1. Transformer（トランスフォーマー）:
   最近の AI の最高峰です。文章を翻訳したり、画像を生成したりするあの技術です。これを粒子の軌跡を見つけることに使おうとしています。
2. FPGA（フィールド・プログラマブル・ゲート・アレイ）:
   これは**「万能な回路基板」です。パソコンの CPU（料理人）のように最初から決まった仕事をするのではなく、「必要な瞬間に、必要な道具（回路）をその場で組み立てて使える」**のが特徴です。
   • メリット: 非常に省エネで、現場（加速器のそば）に設置でき、**「瞬時（リアルタイム）」**に処理できます。

**「AI 探偵を、現場に持ち込める小型の万能工具箱（FPGA）に乗せて、粒子の迷路を瞬時に解き明かそう！」**というのが、このプロジェクトのゴールです。

🔧 挑戦と工夫：どうやって箱に入れるか？

しかし、大きな問題がありました。
最新の AI（Transformer）は、**「巨大な図書館」ほどの大きさがあります。一方、FPGA は「小さな手提げカバン」**ほどの容量しかありません。

「巨大な図書館を、小さなカバンにどうやって詰め込むか？」

著者たちは、以下のような工夫を行いました。

1. 分割して詰め込む（パーティショニング）

図書館全体を一度に持っていけないなら、**「必要なページだけ切り取って」**カバンに入れます。

• フルデプロイ: 図書館全体をカバンに入れる（無理）。
• 部分デプロイ: 重要な章（AI の一部）だけを FPGA で処理し、他の部分は普通のコンピュータに任せる。
  • これなら、どんなに小さなカバン（FPGA）でも、重要な部分だけ高速処理できます。

2. 重さを減らす（量子化）

図書館の本を、**「紙の重さを減らす」**ことで持ち運びやすくします。

• 通常、AI は「10.54321...」のような細かい数字（浮動小数点）を使います。
• FPGA ではこれを「10」や「11」のような簡単な数字（整数）に丸めます。
• 結果: 本が軽くなり、カバンにたくさん入るようになりました。
• 副作用: 丸めすぎると、**「本の情報が少し欠けて、正解率が下がる」**というリスクがありました（特に、計算の途中経過を丸めると精度がガクッと落ちました）。

📊 実験の結果：ZCU102 という「カバン」で試す

彼らは、**「ZCU102」**という特定の FPGA ボードを使って実験しました。

• 容量の限界: このボードには「メモリの棚（BRAM）」が限られています。
• 発見: 1 つの AI の章（エンコーダー層）を入れると、棚の**約 38%**を使ってしまいました。
• 結論: このボードなら、**「最大 4 つの章」**まで入れることができます。それ以上だと、棚が足りなくなります。
  • もし棚が足りない場合は、外付けの大きな倉庫（DDR メモリ）を使えば入りますが、その分、本を取りに行くのに時間（遅延）がかかってしまいます。

💡 この研究のすごいところ（まとめ）

1. 現場でのリアルタイム処理: これまで「後で分析」だった粒子の軌跡を、**「その場で即座に」**AI が判断できるようになりました。
2. 柔軟な導入方法: 巨大な AI を全部乗せる必要はなく、**「必要な部分だけ」**を FPGA に乗せる方法を開発しました。これにより、高性能な機器がなくても導入しやすくなります。
3. 精度と速度のバランス: 「重さを減らす（量子化）」と速度は上がりますが、**「やりすぎると精度が落ちる」**という重要な教訓も得られました。

🎒 日常への例え

想像してみてください。
**「巨大な図書館（AI）」を、「自転車（FPGA）」に乗せて、「交通渋滞の多い街（粒子加速器の現場）」**を走らせようとしています。

• 図書館全体を積もうとすると、自転車が倒れてしまいます。
• そこで、**「一番重要な辞典（AI の一部）」**だけを積みます。
• 辞典を小さく（量子化）して、より多く積もうとすると、ページが読みづらくなります。
• 著者たちは、**「どのページを積めば、最も速く、かつ正確に目的地（正解）にたどり着けるか」**の最適な荷物の組み方を提案しました。

この技術が実用化されれば、将来の大型実験施設では、**「AI が瞬時に『あの粒子は何か』と教えてくれる」**ようになり、新しい物理法則の発見が加速するかもしれません。

技術概要

TrackCore-F: FPGA 上へのトランスフォーマー型素粒子追跡モデルの展開に関する技術的サマリー

本論文は、高エネルギー物理学（特に LHC における素粒子の軌跡再構築）において、Transformer 機械学習モデルを FPGA（Field-Programmable Gate Array）上で効率的に実行するための手法とツールチェーンの構築を目的とした研究です。以下に、問題定義、手法、主要な貢献、結果、および意義について詳細にまとめます。

1. 問題定義 (Problem)

• 計算リソースと遅延の課題: 高エネルギー物理学の実験（特に HL-LHC 段階）では、データ量が爆発的に増加しており、従来のオフライン処理（ポストモーテム処理）では追いつかない状況にあります。ML による軌跡再構築（トラッキング）は計算性能を向上させますが、GPU 依存が主流であり、オンサイト（実験現場）での低遅延・リアルタイム処理には適さない場合があります。
• FPGA 展開の障壁: FPGA は低遅延かつ高エネルギー効率なハードウェアアクセラレータとして有望ですが、Transformer などの大規模 ML モデルを FPGA に展開する際のツールサポートが限られており、モデルの分割や自動統合の手法が確立されていません。
• リソース制約: FPGA のハードウェアリソース（BRAM, LUT, DSP など）は限られており、大規模な Transformer モデルをそのまま展開することは困難です。

2. 手法 (Methodology)

著者らは、TrackFormers プロジェクトに由来する 2 つの Transformer モデル（EncCla と EncReg）を FPGA 上で推論するための開発フローを提案しました。

対象モデル

1. EncCla (Encoder-Classifier): エンコーダのみのトランスフォーマー。ヒット座標を入力とし、トラックパラメータ空間をビンニングしたクラスラベルを予測する分類モデル（約 150 万パラメータ）。
2. EncReg (Encoder-Regressor): エンコーダのみのトランスフォーマー。クラスラベルに依存せず、直接トラックパラメータを回帰するモデル（約 7.6 万パラメータ）。出力には HDBSCAN によるクラスタリングが必要。

開発フローとツールチェーン

• ハードウェア: ARM Zynq UltraScale+ MPSoC (ZCU102) を使用。
• ツールチェーン:
  1. PyTorch: 事前学習済みモデルの提供。
  2. ONNX: モデルの低レベル操作、入出力形状、重みの可視化および量子化のために使用。
  3. AMD Vitis HLS: C/C++ または OpenCL で記述されたハードウェアを FPGA 向けの RTL として合成。
  4. AMD Vivado: IP の統合、MPSoC 構成、ビットストリームの生成。
  5. PYNQ: Zynq PL 部との相互作用を可能にする Python 環境。
• 展開戦略:
  • 部分展開（Partial Deployment）: モデル全体を FPGA に乗せるのではなく、特定のエンコーダ層（Kernel）のみをハードウェアとして実装し、前後の層は CPU 側で処理する「スライス」方式を採用。
  • 手動パーティショニング: 現在のワークフローでは手動でモデルを分割し、エンコーダ層を単一のモノリシックなカーネルとして HLS で実装しています。
  • 最適化: HLS プラグマを用いて性能を最適化。ONNX 演算ブロックを個別に実装するのではなく、セグメントを統合したカーネル化することでリソース使用量を削減しています。

3. 主要な貢献 (Key Contributions)

• FPGA 向け Transformer 展開ワークフローの確立: 事前学習済み PyTorch モデルから FPGA ビットストリームまでの一貫した開発フロー（PyTorch → ONNX → Vitis HLS → Vivado）を提示しました。
• 部分展開の妥当性の証明: モデル全体ではなく、特定の層のみを FPGA で加速する「部分展開」が、リソース制約のある環境でも実用的であり、アクセシビリティを高める有効な手段であることを示しました。
• 量子化の影響評価: 重みとアクティベーションの量子化（INT8/INT16）がモデル精度に与える影響を定量的に評価しました。
• リソース制約の分析: 特定の FPGA 上で実装可能なエンコーダ層の最大数（本ケースでは 4 層）を、BRAM 使用率に基づいて推定しました。

4. 結果 (Results)

• 量子化による精度への影響:
  • 基底モデルの精度は 0.97 です。
  • 重みとアクティベーションをともに INT8 に量子化すると精度は 0.70 まで低下します。
  • 重要な知見: アクティベーションの量子化が精度低下に与える影響が、重みの量子化よりもはるかに大きいことが判明しました（例：重みのみ INT8 でアクティベーションは INT16 の場合、精度は 0.90 を維持）。
• ハードウェアリソース使用量 (ZCU102 上での単一エンコーダ層):
  • BRAM: 38.04% 使用（347/912）。現在のボトルネック。
  • LUT: 24.01% 使用。
  • DSP: 2.66% 使用。
  • FF: 13.53% 使用。
  • 結論: BRAM が主要な制限要因ですが、DDR メモリへの依存を増やすことで削減可能です。LUT 使用率を考慮すると、この FPGA 上では最大 4 層のエンコーダを実装可能と推定されます。
• 性能: 完全な最適化やスループットの数値は現時点では「予備的」とされていますが、部分展開アプローチの有効性が確認されました。

5. 意義と今後の展望 (Significance)

• オンサイト処理の実現: GPU に依存せず、FPGA を利用することで、実験現場での低遅延・高エネルギー効率な ML 推論が可能になります。
• 大規模モデルへの対応: 単一の FPGA に収まらない大規模モデルでも、層ごとの分割（パーティショニング）により実用的なハードウェアで実行可能になることを示唆しました。
• 設計指針の提供: 量子化を行う際のトレードオフ（精度 vs リソース）に関する具体的なデータを提供し、今後の FPGA 向け ML 設計における重要な指針となっています。
• ツールチェーンの進化: 現在のツールサポートの限界を克服するための実用的なワークフローを提示し、HLS 実装自体よりも最適化戦略が重要であるという洞察を与えています。

本論文は、高エネルギー物理学における ML 基盤のハードウェア展開において、Transformer アーキテクチャを FPGA 上で実用的に運用するための道筋を示す重要なステップです。