Apple Silicon 上の MLX でオンデバイス音声認識を構築する

OnType のオンデバイス音声認識は、私たちが開発したオープンソースの Swift ライブラリ mlx-swift-asrによって支えられています。これは Qwen3-ASR を Apple Silicon 上でネイティブに動かすためのものです。 Python ランタイムなし、サブプロセス橋渡しなし、クラウドなし。 Swift → MLX → Neural Engine だけです。

この記事では、推論パイプラインの構造、即時に感じられる体験を支える レイテンシ最適化、そして実際にぶつかった見えにくい落とし穴を説明します。

なぜ Qwen3-ASR で、なぜ MLX なのか

Qwen3-ASR は Alibaba の Qwen チームによる音声認識モデルです。 私たちは 0.6B パラメータ版を 6-bit 量子化して出荷しており、 ディスク上では約 400MB です。Whisper 系の音声エンコーダ （Conv2d → Transformer）から、Grouped Query Attention、RoPE、 SwiGLU を使う Qwen3 テキストデコーダへつながる構成になっています。

Apple の MLX は、Apple Silicon 上での機械学習推論のために作られたフレームワークです。 汎用ランタイムに対する大きな利点は、M シリーズチップの ユニファイドメモリ構造を理解していることです。CPU、GPU、Neural Engine が同じメモリプールを共有するため、ホストメモリとデバイスメモリの 間でデータをコピーする必要がありません。音声が継続的に流れ、 中間表現をすぐ消費するリアルタイム ASR パイプラインでは、 これがレイテンシの大きな削減になります。

しかも私たちはネイティブ Swift バインディングである mlx-swift を使っているので、経路全体が Apple のエコシステム内に残ります。 Swift → MLX → Metal → Neural Engine。Python インタプリタは途中に入りません。

推論パイプライン

mlx-swift-asr における音声からテキストまでの経路は 5 段階です。

Phase 1: Mel spectrogram

16kHz の生音声を、128 ビンの log-mel spectrogram へ変換します。 パラメータは WhisperFeatureExtractor に正確に合わせています。 FFT ウィンドウは 400 サンプル（25ms）、ホップは 160 サンプル（10ms）、 メルフィルタバンク正規化は Slaney 方式です。

Hann window と mel filterbank 行列は起動時に 1 度だけ計算し、 static property としてキャッシュしています。これは小さな最適化 （文字起こし 1 回あたり約 10ms の削減）ですが、すべてのチャンクが 通るストリーミング経路では積み上がります。

見えにくい互換性ポイントが 1 つあります。PyTorch のtorch.stft(center=True) の挙動に合わせるため、 最後の STFT フレームを 1 つ落としています。フレーム数が 1 つずれると、 エンコーダ内で次元不一致が起きます。モデルはクラッシュせず、 ただゴミを出力するだけなので、追うのに非常に時間がかかる種類のバグです。

Phase 2: Audio encoding

Mel spectrogram は音声タワーへ入ります。3 層の stride-2 Conv2d のあとに Transformer encoder が続き、時間軸を約 8 分の 1 に圧縮しつつ、 音声特徴をテキストデコーダの hidden dimension へ射影します。

ここにも細部があります。Conv2d の重みは、モデル読み込み時に PyTorch の OIHW レイアウトから MLX の OHWI レイアウトへ転置する必要があります。 この処理は重みのサニタイズ工程で自動的に行っています。

Phase 3: Prompt construction

エンコード済み音声特徴は、Qwen3 の chat template 形式に従ったテキストプロンプトへ埋め込まれます。プロンプト内の 音声プレースホルダトークンは、cumsum ベースのインデックス方式で 実際の encoder 出力埋め込みへ置き換えられます。構造としては、 マルチモーダル LLM が画像トークンを扱うのと同じ考え方です。

Phase 4: Double-buffering による token generation

最も面白い性能最適化はここです。自己回帰デコードは 1 トークンずつ生成します。素朴にやると、forward pass → token 抽出 → 次の forward pass という直列処理になります。GPU が計算している間 CPU は待ち、CPU がトークンを抜いている間 GPU は遊びます。

私たちは double-buffer の asyncEval パターンで GPU と CPU の仕事を重ねています。

1. 現在のトークンを取り出す 前に、次の forward pass をキューへ積む
2. item() でトークン ID を取り出す。これは GPU sync を強制するが、その時点ですでに GPU は並列で次の logits を計算中
3. 次の logits が必要になる頃には、多くの場合すでに materialize されている

コードの見た目は直感に反します。次ステップ用の input embeddings を準備し、forward pass と asyncEval を積み、 その あとで 現在トークンを抜くからです。 ですが、このパイプライン化によって、本来はデコード時間の大半を占める 同期レイテンシを隠せます。

Phase 5: Token decoding and cleanup

生成された token ID は BPE でデコードし、special token を除去しつつ、検出言語を抽出する output parsing に通します。 その後、クリーンアップ済みテキストはカーソルへ届く前に、私たちの inverse text normalization エンジン を通ります。

Metal warmup: 5 秒のコールドスタート

MLX は実行時に Metal compute kernel をコンパイルします。つまり JIT です。そのため、アプリ起動後最初の文字起こしには約 5 秒の シェーダーコンパイルコストが発生します。2 回目以降は速い。

これはライブラリ全体でもっとも影響の大きい性能上の落とし穴です。 warmup なしでは、最初の本番文字起こしは 8 秒の音声に対して 5〜8 秒かかります（約 0.5x real-time）。warmup ありなら、 以後の文字起こしはすべて 4〜6x real-time で動きます。

この warmup 戦略にも、いくつか見えにくい点があります。

• 無音ではなくノイズを使う。 無音だと mel 値がほぼゼロになり、すべての kernel path を通らない可能性があります。低振幅のランダムノイズなら、 パイプライン全体にわたって多様な計算を確実に踏ませられます。
• 現実的な音声長を使う。 私たちは 8 秒音声で warmup します。これにより batched Conv2d encoder は約 8 チャンクを処理し、実運用に近い batch 次元になります。 2 秒程度だと、最初の実運用文字起こしで、より大きい batch に対する追加の Metal pipeline state compile がまだ走ります。
• 最初の実行では temperature を 0 にしない。 greedy decoding（temperature=0）だと、ノイズ入力に対してモデルが 即座に EOS を出し、トークン生成が 0 で終わることがあります。 すると自己回帰 decode loop の kernel が未コンパイルのまま残ります。 temperature=1.0 にすると、実際にトークン生成が走り、 デコード経路全体を warmup できます。
• 2 回 warmup する。 1 回目は sampling 付きで全 kernel をコンパイルし、2 回目は greedy decoding で fast path を検証します。その後メモリキャッシュは消しますが、 シェーダーキャッシュは残します。

ベンチマーク

対象は M1 Pro 上の 6-bit 量子化 Qwen3-ASR 0.6B です。 以下はモデルロードと Metal warmup を除いた、純粋な推論時間です。

典型的な RTF は 0.15〜0.27（3.7〜6.3x real-time）です。 Python 版 MLX 実装と競争力がありつつ、Python オーバーヘッドなしの ネイティブ Swift ライブラリとして動作します。

OnType のリアルタイム音声入力という用途では、これはつまり、 音声が到着するより速く ASR エンジンが処理できるということです。 ユーザーがホットキーを離す時点で、音声の大半はすでに文字起こし済みです。 残る最後のチャンクだけを処理すればよく、それは 300ms 未満で終わります。

ガードレール

痛い目を見ながら追加した安全装置をいくつか挙げます。

• 音声長ベースの token cap。 最大生成長はceil(audioDuration × 20) + 64 token に制限しています。これがないと、病的な入力で EOS を出さないまま 数千 token を生成し続け、GPU が延々と回り続けることがあります。
• 反復検知。 同じ token が 10 回連続したら停止します。 これは、モデルがループにはまる退化出力を捕まえるためです。
• 最小音声長。 FFT window より短い音声 （400 サンプル = 25ms）は即座に空結果で返します。 このチェックがないと、STFT の reflection padding で不正範囲クラッシュが起きます。
• MLX error conversion。 MLX は GPU エラー時の既定動作が fatalError です。そこで重要な経路を withError で包み、Swift の throws へ変換しています。これにより、アプリはクラッシュする代わりに HUD へエラーを表示できます。

オープンソース

mlx-swift-asr は MIT ライセンスで、 github.com/ontypehq/mlx-swift-asrで公開しています。単独の Swift package なので、 Package.swift に追加するだけで、3 行 API でオンデバイス音声認識を使えます。

let stt = try await Qwen3ASRSTT.loadWithWarmup(from: modelDirectory)
let result = try await stt.transcribe(file: audioURL)
print(result.text)  // "Hello, world."
print(result.rtf)   // 0.17 (5.8x real-time)

Apple Silicon 向けに音声認識が必要なもの （音声入力、文字起こし、アクセシビリティツールなど）を作っているなら、 ぜひ試してください。ベンチマークは swift test --filter Benchmark で再現できます。

OnType を試す — mlx-swift-asr の実動作を体験できます。 キーを押して話し、離すだけで入力。オンデバイス、リアルタイム、 そしてプライベートです。