WhisperKitのメモ

WhisperのCore ML版であるWhisperKitに関するメモ。

WhisperKitによるiOSオンデバイス音声認識、Releaseビルドで試すとさらに爆速でした。15分の音声ファイルの書き起こしが51秒で完了。baseモデル利用、デバイスはiPhone 15 Pro。 pic.twitter.com/He3a0FQntn

— 堤修一 / Shuichi Tsutsumi (@shu223) March 27, 2024

今日のLTでWhisperKitについて話すのだけど、5分だと話しきれないこともいろいろあるのでこちらに書いていく。

whiper.cppとの違い

公式ブログの記事の "Accelerating the Audio Encoder" という節で、WhisperKitの系譜のようなものについて言及があり、その中でwhisper.cppとの関係性についても触れられている。

The audio encoder mentioned in Step 3 was previously optimized by the whisper.cpp project using Core ML following Deploying Transformers on the Apple Neural Engine. Simply put, this work introduces a set of Neural Engine compiler hints as PyTorch code that translate to a high-performing model when converted to Core ML. Given that the Neural Engine compiler is Apple’s private framework, this is a form of black-box optimization to nudge the compiler for increased hardware utilization.
（ステップ3で述べたオーディオエンコーダは、whisper.cppプロジェクトによって、Deploying Transformers on the Apple Neural Engineの後にCore MLを使って最適化されました。簡単に言うと、この仕事は、PyTorchコードとしてNeural Engineコンパイラのヒント を導入し、Core MLに変換したときに高性能なモデルに変換します。Neural EngineコンパイラがAppleのプライベートフレームワークであることを考えると、これはブラックボックス最適化の一形態であり、ハードウェアの利用率を高めるためにコンパイラを誘導する。）

This is a great starting point for accelerating Whisper’s audio encoder and we are happy to see that our previous work is being adopted in various open-source projects.

これは、Whisperのオーディオエンコーダを高速化するための素晴らしい出発点であり、私たちのこれまでの研究が様々なオープンソースプロジェクトで採用されていることを嬉しく思います。

In WhisperKit, we further improved these compiler hints which yielded an additional 1.86x-2.85x speedup across iPhone 12 through 15 compared to the previous best implementation:

WhisperKit_では、これらのコンパイラのヒントをさらに改善し、iPhone 12からiPhone 15にかけて、以前の最高の実装と比較して、さらに1.86倍から2.85倍のスピードアップを実現しました：

ざっくり要約：

• whisper.cpp は、apple/ml-ane-transformers の知見をベースに、ANEの性能を引き出すようにCore MLモデルに変換されている

• WhisperKit はさらにその最適化を推し進めたもので、さらに1.86倍〜2.85倍の高速化に成功した

argmax社と apple/ml-ane-transformers

上述のargmax社のブログ記事の中で参照されている「Deploying Transformers on the Apple Neural Engine」というLinkedInの記事（argmax社ファウンダーｂのAtila Orhon氏によるもの）を読むと、

Atila Orhon on LinkedIn: Deploying Transformers on the Apple Neural Engine | 12 comments

As part of hashtag#wwdc2022, we are open-sourcing a reference implementation of the Transformer architecture optimized for Apple Neural Engine (ANE)! We also just published an accompanying research article to describe the principles behind the optimizations and empower developers to make informed decisions on ANE deployment of their ML models.
（#wwwdc2022の一環として、我々はApple Neural Engine（ANE）用に最適化されたTransformerアーキテクチャのリファレンス実装をオープンソース化しています！また、最適化の背後にある原理を説明し、開発者がMLモデルのANEデプロイメントについて十分な情報を得た上で意思決定できるようにするために、付随する研究論文を発表しました。）

In our research article case study, the Hugging Face distilbert model is up to 10x faster with 14x lower peak memory consumption after our optimizations are applied on it while consuming as low as 0.07W of power.
（私たちの研究論文のケーススタディでは、ハギング・フェイス・ディスティルベルト・モデルに私たちの最適化を適用したところ、消費電力が0.07Wと低い一方で、最大10倍高速化し、ピーク・メモリ消費量が14倍減少した。）

This work would not have been possible without extensive collaboration between hardware and software teams, enabling us to optimize across the stack.
（この作業は、ハードウェア・チームとソフトウェア・チームの広範な協力なしには不可能であり、スタック全体の最適化を可能にした。)

Our applied R&D team is leveraging this implementation in production for technologies such as HyperDETR which we described in a previous research article.
（私たちの応用研究開発チームは、以前の研究論文で紹介したHyperDETRなどの技術に、この実装を実生産で活用している。）

Code: https://lnkd.in/gYBJrGKy

Deploying Transformers on the Apple Neural Engine (Research Article): https://lnkd.in/guWGg9-J

HyperDETR (Research Article): https://lnkd.in/gWJkWvFd

とあり、Appleの apple/ml-ane-transformers やApple ML Blogの記事がリンクされている。

え、argmaxの人がこのAppleの仕事やったの？と当該リポジトリのコミット履歴をみると、そういうことのようだ。

つまりWhisperKitは然るべき実績がある人によってカリカリ最適化されたもので、速いのにはちゃんとした裏付けがある。

WhisperKitを出してるargmax社、そもそもこのTransformerをANE最適化するAppleのリポジトリの中の人の会社だった。だからカリカリに最適化できるのか。

apple/ml-ane-transformershttps://t.co/C00PVfgaGg

— 堤修一 / Shuichi Tsutsumi (@shu223) March 27, 2024

WhisperKitのREADMEに書いてあったこと

前述の公式ブログ記事に比べ、こちらは使い方の説明に特化したREADMEだった。

WhisperKit is a Swift package that integrates OpenAI's popular Whisper speech recognition model with Apple's CoreML framework for efficient, local inference on Apple devices.

• 最小実装

let pipe = try? await WhisperKit()
let transcription = try? await pipe!.transcribe(audioPath: "path/to/your/audio.{wav,mp3,m4a,flac}")?.text
print(transcription)

• モデルの選択方法

• モデルの生成方法

• Swift CLI実装もあるよ

• 書いてない：

  • 詳しい中身の解説

  • whisper.cppやMLXとの違い

Hugging Faceの whisperkit-coreml に書いてあったこと

argmaxinc/whisperkit-coreml · Hugging Face

各モデルの精度とファイルサイズ

まず、各モデルの精度とファイルサイズの表が掲載されている。

• Short-form Audio と Long-Form Audio に対する精度（WER）で表が分かれている

• 比較用に OpenAI の Whisper API large-v2 で同様のテストを行った結果も掲載されている

それについて詳細な解説が続く（以下）

We believe that rigorously measuring the quality of inference is necessary for developers and enterprises to make informed decisions when opting to use optimized or compressed variants of any machine learning model in production. To contextualize WhisperKit, we take the following Whisper implementations and benchmark them using a consistent evaluation harness:
（推論の質を厳密に測定することは、開発者や企業が最適化または圧縮された機械学習モデルを実運用で使用する際に、十分な情報に基づいた意思決定を行うために必要であると考えています。WhisperKitの文脈を理解するために、以下のWhisper実装を取り上げ、一貫した評価ハーネスを使用してベンチマークを行います：）

Server-side:
・WhisperOpenAIAPI: OpenAI's Whisper API
($0.36 per hour of audio as of 02/29/24, 25MB file size limit per request)

On-device:WhisperKit:
・Argmax's implementation [Eval Harness] [Repo]
・whisper.cpp: A C++ implementation form ggerganov [Eval Harness] [Repo]
・WhisperMLX: A Python implementation from Apple MLX [Eval Harness] [Repo]
(All on-device implementations are available for free under MIT license as of 03/19/2024)

ここでwhisper.cppやMLXについて言及されているが、表にはこれらの検証結果が載ってないので、どういう話かわからなかった…

QoI

表にWERと併記されている指標 QoI についても以下のように解説がある：

WhisperOpenAIAPI sets the reference and we assume that it is using the equivalent of openai/whisper-large-v2 in float16 precision along with additional undisclosed optimizations from OpenAI. In all measurements, we care primarily about per-example no-regressions (quantified as qoi below) which is a stricter metric compared to dataset average Word Error RATE (WER). A 100% qoi preserves perfect backwards-compatibility on the test distribution and avoids "perceived regressions", the phenomenon where per-example known behavior changes after a code/model update and causes divergence in downstream code or breaks the user experience itself (even if dataset averages might stay flat across updates). Pseudocode for qoi:
（WhisperOpenAIAPIはリファレンスを設定し、OpenAIからの非公開の追加最適化とともに、float16精度でopenai/whisper-large-v2と同等のものを使用していると仮定します。すべての測定において、私たちは主に、データセット平均のワードエラー率（WER）と比べてより厳しい指標である、サンプルごとの無回帰（以下qoiと定量化）を気にしています。qoiが100%であれば、テスト分布の完全な後方互換性が保たれ、"perceived regressions"（コード/モデルの更新後にサンプルごとに既知の動作が変化し、ダウンストリームのコードに乖離を引き起こしたり、ユーザーエクスペリエンスそのものを壊してしまう現象）を回避することができます（たとえデータセット平均が更新をまたいで横ばいであったとしても）。qoiのシュードコード）

qoi = []
for example in dataset:
    no_regression = wer(optimized_model(example)) <= wer(reference_model(example))
    qoi.append(no_regression)
qoi = (sum(qoi) / len(qoi)) * 100.

Note that the ordering of models with respect to WER does not necessarily match the ordering with respect to QoI. This is because the reference model gets assigned a QoI of 100% by definition. Any per-example regression by other implementations get penalized while per-example improvements are not rewarded. QoI (higher is better) matters where the production behavior is established by the reference results and the goal is to not regress when switching to an optimized or compressed model. On the other hand, WER (lower is better) matters when there is no established production behavior and one is picking the best quality versus model size trade off point.
WERに関するモデルの順序は、QoIに関する順序と必ずしも一致しないことに注意してください。これは、参照モデルには定義上100%のQoIが割り当てられるからである。他の実装による例ごとの回帰はペナルティを受けるが、例ごとの改善は報われない。QoI（高ければ高いほど良い）は、本番の動作が参照結果によって確立され、最適化または圧縮されたモデルに切り替えたときに後退しないことを目標とする場合に重要である。一方、WER（低ければ低いほど良い）は、確立されたプロダクション動作がなく、最適な品質対モデルサイズのトレードオフ・ポイントを選択する場合に重要です。

We anticipate developers that use Whisper (or similar models) in production to have their own Quality Assurance test sets and whisperkittools offers the tooling necessary to run the same measurements on such custom test sets, please see the Model Evaluation on Custom Dataset for details.
Whisper（または類似のモデル）を本番で使用する開発者は、独自の品質保証テストセットを持っていることが予想され、whisperkittools はそのようなカスタムテストセットで同じ測定を実行するために必要なツールを提供します。

Why are there so many Whisper versions?

WhisperKit is an SDK for building speech-to-text features in apps across a wide range of Apple devices. We are working towards abstracting away the model versioning from the developer so WhisperKit "just works" by deploying the highest-quality model version that a particular device can execute. In the interim, we leave the choice to the developer by providing quality and size trade-offs.
（WhisperKitは、様々なAppleデバイスのアプリに音声合成機能を組み込むためのSDKです。WhisperKitは、特定のデバイスが実行可能な最高品質のモデルバージョンをデプロイすることで、「単に動作する」ように、開発者からモデルのバージョン管理を抽象化することを目指しています。暫定的には、品質とサイズのトレードオフを提供することで、開発者に選択を委ねています。）

（DeepLの翻訳が相当イマイチだが…）つまり要はたくさんのモデルがあるのは品質とサイズのトレードオフについて開発者に選択肢を与えてるってことをいいたいんだと思う。

最後に、検証結果の再現方法についても言及されている。あなたの手元のApple Sillicon Macで同じことができますよ、と。

WhisperKitサンプルのソースコードを読む（別記事）

別記事に書きました