TensorFlow を使用した ML モデルの実装と最適化: 1 秒あたり 3 億回の予測

[[425184]]

TensorFlow は最も広く使用されている機械学習フレームワークの 1 つであり、研究を加速し、新しいモデルの作成時間を短縮します。ネイティブプログラマティック DSP 企業 Zemanta のデータサイエンス ディレクター Davorin Kopič 氏とエンジニア Jan Hartman 氏は、論文の中で、オンライン広告エコシステム内の大規模な機械学習モデルを TensorFlow フレームワークに変換するプロセスと、TensorFlow フレームワークに実装された機械学習モデルを 1 秒あたり 3 億回以上の予測に拡張するプロセスを実演しました。したがって、この研究の主な焦点は、TF にモデルを実装し、さまざまな最適化手法を使用して低レイテンシで効率的に提供することです。

論文アドレス: https://arxiv.org/abs/2109.09541

この研究で使用されるケースは、オンライン広告のクリック予測です。 RTB（リアルタイム入札）では、複数の DSP（入札者）が、Web ページの読み込み中にリアルタイムで入札してオンライン広告スペースを競います。広告スペースは広告インプレッションごとに販売されるため、仮想広告スペースを市場価格で販売することが可能です。 RTB では、機械学習を使用することで、広告主がクリックスルー率 (CTR) などの KPI を最大化することもできます。広告のクリックスルー率を予測することは RTB の中心的な課題の 1 つであり、優れたクリック予測モデルを持つことは非常に重要です。

Golang で実装されたカスタム ロジスティック回帰および因数分解マシン (FM) に基づくモデルは表現力が限られており、すべての学習手順を手動で実装する必要があるため、実験が遅くなり、モデルの予測パフォーマンスが制限されます。そのため、研究者たちは TensorFlow フレームワークを採用し、既存のモデルをより表現力豊かなモデルに置き換えることにしました。

課題

広告入札の特定のユースケースに基づいて、調査ではいくつかの課題に遭遇しましたが、これは実装、サービス、最適化の 3 つの側面に分けられます。

一方で、各マシンに 1 つ以上の最高級 GPU を装備すると、コストがかかりすぎます。また、GPU のクラスターが小さいマシンでは、研究をサービスベースのアーキテクチャに移行せざるを得なくなります。どちらのオプションも特に望ましいものではなく、また、私たちのモデルはコンピューター ビジョンや自然言語処理などの他のディープラーニング ドメインの SOTA モデルと比較して比較的小さいため、この研究では、実稼働環境での推論に GPU を使用しません。また、この研究のモデルはスパース重みを使用しているため、その使用例は GPU ワークロードには適していません。

成し遂げる

TF モデルの効率的なトレーニング ループを実装するために、この研究では複数の方法を実装してテストしました。 TF における高スループットのオンライン トレーニングとサービスのケース スタディは少なく、ドキュメントも十分に具体的ではないことが多く、研究者はソース コードを読み、プロトタイプをベンチマークして実装の落とし穴を発見する必要があります。

TF は、SOTA アルゴリズム実装を備えた巨大なエコシステムと多数のライブラリを提供します。豊富な機能を備えた既存の実装を選択するのは簡単ですが、研究者はこれらの実装のほとんどが最適化されていないことに気づいたため、アルゴリズムを独自に実装することにしました。 TF にはさまざまな抽象化レベルの API がありますが、一部の API は使いやすいものの、効率の悪い低レベルの操作であることがよくあります (Estimator API など)。研究者が最終的に Keras3 を選択したのは、それが基礎となる TF 操作を薄くラッパー化したものであり、理解しやすいだけでなく高いパフォーマンスを発揮するという理由からです。 TF は非常に機能とリソースが豊富なライブラリであるため、この調査では機械学習パイプラインのどの程度が TF に実装されるかについても考慮する必要がありました。研究者たちは、機能の変換と相互作用を一時的に脇に置いて、学習アルゴリズムのみを実装することを選択しました。これらは交換可能な最小の部品ですが、改善の可能性が最も大きいです。

Golang TF ラッパーは予測のみをサポートしているため、トレーニング ループは Python で実装する必要があります。スクリプトは、標準入力をサブプロセスとして渡すことで、Golang データ パイプラインに接続します。データは解析を必要としない効率的なバイナリ形式で送信されるため、この方法は CSV 形式よりも 25% 高速になります。次に、データを待機している間にモデルがアイドル状態にならないように、データがバックグラウンド スレッドで読み取られます。これに基づいて、この研究ではトレーニング パイプライン全体にわたって高いスループットを達成しました。効率的な入力と出力も低遅延予測の鍵であることがわかりました。この研究では、すべての入力機能を 1 つのテンソルに連結することで、入力データのシリアル化とコピーにかかる時間を大幅に短縮しました。

仕える

研究者らは、計算集約型のニューラル ネットワークのため、Golang TF デコレータを使用すると DeepFM モデルの CPU 使用率が大幅に高くなることを発見しました。パフォーマンスは大幅に向上しますが、このアプローチをトラフィックの 100% に拡張すると、多大なハードウェア コストが発生します。現在の世界的なチップ不足を考えると、これはコストが困難で高価になることを意味します。

明らかに、計算コストを削減する必要があります。ただし、ニューラル ネットワーク モデルのサイズを縮小すると、モデルの予測パフォーマンスも低下します。研究者らはTFをさらに研究した結果、バッチ計算時に例の数を増やすと計算効率が大幅に向上することを発見した。この低い線形増加は、TF コードが高度にベクトル化されており、TF では計算呼び出しごとにオーバーヘッドが発生し、それがバッチで償却されるためです。これを念頭に置いて、コンピューティング呼び出しの数を減らしたい場合は、多数のリクエストを 1 つのコンピューティングに連結する必要があります。

研究者らは、ネットワーク呼び出しを回避するために、実行中の単一の入札インスタンスに完全に含まれる自動バッチ処理システムを構築しました。各インスタンスは 1 秒あたり数千の受信リクエストを受信するため、多数のリクエストからの計算が確実に接続され、より大きなバッチが作成されます。研究者たちは、着信リクエストからデータを受信し、バッチを作成し、バッチが完了した後に計算を初期化する多数のバッチ スレッドを通じてこれを実現しました。計算中は、この時間枠内にバッチが満たされない可能性があるため、タイムアウトを回避するために数ミリ秒ごとに初期化されます。この実装は高度に最適化されており、TF 計算の CPU 使用率を半分にしながら、計算呼び出しの数を 5 分の 1 に削減します。

まれにバッチ プロセッサ スレッドが CPU 時間を取得できず、要求タイムアウトが発生することがありますが、これは要求の 0.01% 未満で発生します。研究者らは、平均待ち時間がわずかに増加していること（平均で約 5 ミリ秒）を観察しましたが、トラフィックがピークになる時間帯にはさらに増加する可能性があります。そのため、レイテンシの安定性を確保するために、SLA（サービスレベル契約）と適切な監視対策を実施しました。タイムアウトの割合が大幅に増加しなかったことを考えると、これらの方法は依然として非常に効果的であり、この TF サービス メカニズムの中核となっています。

[[425186]]

共著者 ダヴォリン・コピッチ

最適化

TF で実装したモデルは、当初はカスタム FM よりもかなり遅いものでした。高速化の余地を見つけるために、組み込みの TF プロファイラーを多用して、実行に最も時間のかかる操作を見つけ、可能な限り改善しました。最も一般的なのは、さまざまな冗長な再形成または変換操作です。興味深い発見の 1 つは、操作数の差は小さいにもかかわらず、Adam オプティマイザーは Adagrad よりも大幅に遅い (約 50%) ことです。プロファイラーは、スパース重みの勾配更新に多くの計算時間が必要であることを示しています。これは、モデルの重みがスパース (特徴はほとんどがカテゴリであるため非常にスパース) であり、オプティマイザーがこの事実を考慮していないためです。

Adam を Adagrad に置き換えると、ディープ モデルのパフォーマンスが大幅に低下するため、研究者は他の解決策を模索しました。Lazy Adam オプティマイザーへの切り替えは、スパースな重みを非常に効率的に処理できるため、非常に効果的であることが判明しました。結果を見ると、全体的なトレーニング速度が 40% 以上加速され、Adagrad に近い値になっていることがわかります。

適応型オプティマイザー（Adam など）が使用され、重みモーメントと分散も保存する必要があるため、パラメータごとに 3 つの値が保存され、保存されるモデル サイズが 3 倍に増加します。ただし、これらの値は実際には予測には使用されず、トレーニングにのみ使用されます。研究者たちはこれを利用して、モデルからこれらの値を削除する最適化プロセスを構築し、データ量を 66% 削減し、メモリ使用量とコストを下げました。