【ITニュース解説】EdgeBERT: I Built My Own Neural Network Inference Engine in Rust

システムエンジニアを目指す皆さんに、画期的な技術「EdgeBERT」について解説しよう。この技術は、私たちが普段利用するアプリケーションの裏側で、どのようにテキストの意味を理解し、より賢い機能を実現しているのか、そしてその実現がいかに大変な道のりであったかを示している。

まず、私たちが直面した課題は「セマンティック検索」の実現だった。これは、例えば検索窓に「医者」と入力したときに、単に「医者」という文字列だけでなく、「医師」や「医療従事者」といった、意味的に関連性の高い言葉を含む情報も見つけ出せるようにする技術だ。この機能を実装するためには、テキストデータを単なる文字列としてではなく、その「意味」を数値の並び（ベクトル）として表現する必要がある。つまり、「医者」と「医師」は意味が似ているため、それらを数値化したベクトルも互いに近い値になるように変換するのだ。これにより、数値同士の距離を計算することで、意味的な関連性を判断できるようになる。

このテキストからベクトルを生成する技術は「埋め込み（Embedding）」と呼ばれ、特にBERTのような大規模言語モデルがよく使われる。しかし、この埋め込み生成をアプリケーションに組み込むことは、これまで多くの課題を伴ってきた。

一般的な解決策として、Pythonエコシステムが提供するSentence-Transformersのようなライブラリを利用する方法がある。これは非常に強力で便利なツールだが、その欠点は「重さ」にある。埋め込み生成のためだけに、必要なライブラリ（PyTorch、トークナイザー、モデルの重みなど）をインストールすると、仮想環境全体で6.8GBものディスク容量を消費してしまうことがあった。これでは、リソースが限られた環境や、デプロイサイズを最小限に抑えたい場合には、大きな障壁となる。

次に、Pythonを使わずにRustで実装しようと試みた際も、新たな困難が浮上した。ONNX Runtime（ort）という、異なるフレームワークで学習された機械学習モデルを実行できるライブラリをRustから利用する方法だ。これを使えば、Python環境が不要になるかと思われたが、実際にはRustのコードが複雑化し、基本的な埋め込み機能を実現するだけで150行ものコードが必要になった。さらに深刻だったのは、その依存関係の多さだ。ortを組み込むと、80以上のRustクレート（ライブラリ）がプロジェクトに追加され、結果としてコンパイルされたバイナリサイズは350MBにも膨れ上がった。

このバイナリの肥大化だけでなく、システムライブラリとの衝突も大きな問題となった。ortは、OpenSSLのようなシステムにインストールされているC++ライブラリに依存することがある。私の環境では、ortが要求するOpenSSLのバージョン（3.3）と、システムにインストールされているバージョン（1.1）が異なっていた。システムに古いOpenSSLバージョンがインストールされているのは、他の多くのシステムコンポーネントがそれに依存しているためで、安易にアップグレードすると他のシステム機能が動作しなくなるリスクがあった。このように、軽量なオフラインAIソリューションを目指していたにもかかわらず、たった一つの依存関係がプロジェクト全体を頓挫させるほどの大きな課題となってしまったのだ。

そこで私は、これらの問題を根本的に解決するため、Rustで独自のニューラルネットワーク推論エンジン「EdgeBERT」を構築する決断をした。目標は、PythonのSentence-Transformersのようにシンプルで使いやすいAPIを提供しつつ、軽量で、依存関係の煩わしさがないソリューションを作ることだった。

その結果は目覚ましいものだった。EdgeBERTはわずか5MBのバイナリサイズで、メモリ消費量も200MB程度に抑えられた。依存関係はほとんどなく、前述のようなシステムライブラリとの競合問題も解消された。そして何よりも重要なのは、その精度がPythonのSentence-Transformersとほぼ同じ（0.9997のコサイン類似度）であることだ。これは、浮動小数点演算のわずかな丸め誤差を除けば、実質的に同じ結果が得られることを意味する。

パフォーマンス面でも、EdgeBERTは優れている。特にシングルスレッド環境では、既存のSentence-Transformersに比べて最大3倍高速に動作することが確認された。これは、MiniLMのような比較的小さなモデルの場合、複数のスレッドで処理を並列化する際の「スレッド間調整のオーバーヘッド」が、並列化によるメリットを上回ることがあるためだ。EdgeBERTはこのオーバーヘッドを最小限に抑えることで、軽量アプリケーションやエッジデバイスにおいて特に高いパフォーマンスを発揮する。また、メモリ使用量も非常に効率的で、Pythonベースのソリューションに見られるような初期の大規模なメモリ割り当てスパイクがなく、安定したメモリフットプリントを維持する。

EdgeBERTの最大の利点は、これらの機能をわずか5MBのRustバイナリで提供できる点にある。これにより、Python環境を一切必要とせず、スマート検索、レコメンデーション、重複検出、コンテンツモデレーション、RAG（Retrieval Augmented Generation）のようなAI機能を、メモリやストレージ容量に制約のあるエッジデバイス、あるいはWebブラウザ上でも簡単に実現できるようになる。実際、EdgeBERTはWebAssembly（WASM）形式にもコンパイル可能で、わずか429KBのWASMバイナリと30MBのモデルウェイトでブラウザ内で動作する。ただし、WASM対応のためには、一部のC言語に依存するトークナイザーをRustでゼロから再実装する必要があった。

EdgeBERTの内部動作は、基本的にBERTモデルの推論プロセスをRustで忠実に実装したものだ。これは、テキストの「トークン化」（単語をIDに変換）、そこから「埋め込み」を生成し、複数の「トランスフォーマー層」（自己アテンションとフィードフォワードの繰り返し）を通して意味表現を洗練させ、最終的に「プーリング」（各トークンの埋め込みを統合して文全体の埋め込みを生成）という一連の行列演算で構成されている。コア実装は約500行のRustコードにまとめられており、複雑な外部依存なしに、このアルゴリズムを効率的に実行している。

今後の展望としては、GPUサポートの追加、より多くのモデルアーキテクチャへの対応、そしてモデルの「量子化」（モデルサイズをさらに小さくし、CPUやマイクロコントローラーでの性能を向上させる技術）の実装が挙げられている。

このEdgeBERTの登場は、システムエンジニアを目指す皆さんにとって、ソフトウェア開発における「軽量性」「依存関係管理の重要性」「パフォーマンス最適化」がいかに重要であるかを示す良い事例となるだろう。重厚な既存ソリューションに代わる、効率的でスリムな選択肢が求められる現代において、このような「ゼロから構築する」アプローチが新たな可能性を切り開くことを示している。EdgeBERTは、まさに「必要は発明の母」を体現したプロジェクトであり、より多くの開発者がこの恩恵を受けられるようになることが期待される。