量子位出典公開日

KV Cacheが定数に圧縮!OCRオープンソースの新王者:ページ数に上限なし、キャッシュは96%削減

听雨 凹非寺より 量子位 | 公衆アカウント QbitAI

听雨 凹非寺より
量子位 | 公衆アカウント QbitAI

百度のOCRがまたアップグレードした。文心Unlimited OCR 2が正式リリースされ、オープンソース化!

今回は、一気に小説全巻を読み切ることに直接挑戦した。

『ハリー・ポッター』全7部、『賢者の石』から『死の秘宝』まで、連続で認識・転写し、3600+ページを最後まで読み通した。

『三体』全集に切り替えても、同様に一気に読み切ることができ、3冊で1200+ページ、数十万字を途切れることなく読み進めた。

ただし、これほど多くのページを連続して読むには、モデルはまず一つの関門をクリアしなければならない:VRAM(ビデオメモリ)。

文書が長くなるほど、推論時に保持すべき中間情報(KVキャッシュ)は増えていく——これは、モデルが読みながら取るメモのようなもので、ページ数が多いほどメモは分厚くなる。メモが分厚いと、場所を取るだけでなく、めくるのも遅くなる。

技術レポートの百ページ級シングルシーケンステストにおいて、文心Unlimited OCR 2は一组の比較を示した:

ピーク時のKVキャッシュ使用量は、標準デコード(すべてのページの視覚キャッシュを同時にVRAMに保持する方式)の47GBから1.77GBに低下し、一気に約96%削減された。

47GBとはどのくらいの規模か?多くのノートパソコンの本体メモリよりも大きい。しかもこれはKVキャッシュだけで、モデルの重みなどのその他のオーバーヘッドは別途計上される。

その前世代にあたるUnlimited OCRは、多ページ文書の連続解析能力によって、すでにオープンソースコミュニティで相当な人気を集めていた:累計ダウンロード数は1000万を超え、GitHubでは約2.86万Starを獲得している。

コミュニティによるオンラインDemoの作成から、SGlang推論のサポートまで、それを取り巻く適合作業も継続的に進められている。

第2世代でのアップグレードの重点は、より長い文書を連続処理すると同時に、計算とキャッシュの負荷を下げることにある。

このため、文心Unlimited OCR 2はスパース参照スライディングウィンドウ注意機構SR-SWAを提案し、モデルが必要に応じて関連ページと視覚情報を選択できるようにし、重複計算を削減した。

継続デコード機構と組み合わせることで、モデルは超長文書を連続処理でき、1ステップあたりのデコードコストも文書が長くなるにつれて増えることはない。

現在、文心Unlimited OCR 2はHugging FaceとGitHubで正式にオープンソース化されている。

オープンソースのアドレス:
*GitHub:https://github.com/baidu/Unlimited-OCR*
*HuggingFace:https://huggingface.co/baidu/Unlimited-OCR-2*
コア技術の解説:SR-SWAはどう計算を減らし、チャンクデコードはどう保存を減らすのか
数百万字、3600+ページに及ぶ『ハリー・ポッター』全集を連続で読み切った文心Unlimited OCR 2は、一体どのようにしてそれを実現したのか?
それが参考にしているのは、実は人間が本を写すときのやり方だ。
人が手で本を写すときの流れはどうか。一節を見て、一節を写し、目は主に今写しているページに向いている。前後を照らし合わせる必要がある場所に遭遇したときだけ、関連するページに戻ってめくる。
本を写すという行為について、人類はモデルより早く理解していた。誰も『論語』を写しながら、頭の中で『論語』の全ページを同時に広げて一緒に見たりはしない。
しかし前世代のUnlimited OCRはそうしていなかった。R-SWAによって推論段階のKVキャッシュを大幅に削減してはいたものの、再び見返す過程では、すべての入力ページの視覚情報が依然としてアテンション計算に参加する必要があった。
文書が長くなるほど計算する情報は増えるが、そのうち相当な部分は、現在認識している数行の文字とはまったく関係がない——典型的な「お付き合い」状態だ。

△SR-SWAは人間が本を写すときの局所的な集中方式を参考にしている
人間が本を写す流れに沿って、文心Unlimited OCR 2は次を導入した。疎参照スライディングウィンドウアテンション機構(SR-SWA)。
その選別は2段階に分かれている。まず参照が必要なページをふるい出し、次にふるい出されたページの中から実際に関連する局所情報を選び出し、無関係な計算を締め出す。
アーキテクチャ全体から見ると、文心Unlimited OCR 2は前世代の基本設計を受け継いでいる。

△文心Unlimited OCR 2の全体アーキテクチャ
書籍ページの画像はまずDeepEncoderに入り、モデルが処理できる視覚Tokenへと変換される。これは、ページの内容やレイアウトの特徴を記録したデータカードの集まりと理解できる。
次に、言語モデルがこれらのカードに基づいて、一歩一歩認識結果を生成していく。
重要な変化は言語モデルの内部で起きている。前世代で採用されていた参照スライディングウィンドウアテンション機構(R-SWA)が、この世代でSR-SWAにアップグレードされた。
アーキテクチャ図の左下に示されているのが、SR-SWAの2段階の選別プロセスだ。
第一層はPage Gateで、「どのページを見るべきか」を担当する。
モデルが現在解析している位置に基づいて、候補の視覚情報を関連ページの範囲にまとめ、無関係なページが混ざり込まないようにする。
本を写す場面に置き換えると、あるページを読むときに、まずそのページに関連するいくつかのページを見つけるのであって、最初から本全体をめくるのではない、という具合だ。
第二層はToken Sieveで、「そのページのどの行を見るべきか」を担当する。
ページを正しく選ぶことは第一歩にすぎない。たとえこれらのページの中でも、すべての視覚Tokenを計算に参加させる必要はない。Token Sieveはさらに選別を行い、ページの中から現在の内容により関連する視覚情報ブロックを選び出し、長距離の視覚依存関係を捉える。
アーキテクチャ図の右側にはこのプロセスが展開されている。まず各情報ブロックの要約表現で関連性を算出し、次にTopKによってスコアが最も高い、数量固定の情報ブロックを選び出す。
選び出した後も、モデルはこれらのブロックの中の具体的な視覚情報を見るのであって、要約だけを見るわけではない。
2層の機構が連携することで、モデルは長文書の中から関連ページを見つけられるうえ、各ステップで処理すべき視覚情報を減らすことができ、長文書解析の効率が向上する。

△Token Sieveの詳細
ふるい出された視覚情報は、図中のLocal Window(ローカルスライディングウィンドウ)とともに解析に参加する。
このウィンドウのサイズは固定で、直近の文脈を格納しており、生成プロセスに沿って前へとスライドし続け、モデルが直前の内容に沿って書き進められるようにする。
情報ブロックの数は固定、ローカルウィンドウのサイズも固定であり、各ステップで参照する情報量が枠に収められている。文書がどれだけ長くなっても、このステップの作業量は同じままだ。「本が厚くなるほど、一文字書くのが大変になる」という事態は起きない。
SR-SWAが担当するのは各ステップでの計算量を減らすこと。一方、メモリにどれだけの情報を残すべきかは、別の問題だ。
各ページの画像はエンコードされると、一組の視覚Tokenに変換されます。モデルは解析結果を生成する際にこれらを参照するため、対応するKVキャッシュを保存し、以降の計算で再利用できるようにします。
文書全体を一度に読み込み、すべてのページの視覚キャッシュをVRAMに残したままにすると、ページ数が多いほど占有面積は大きくなります。
さらに厄介なのは、これらのページが現在の計算に一切関与していなくても、キャッシュが解放されない限り、そのままVRAMを占有し続けることです。これは、書物を写す際に各ページの草稿をすべて机の上に広げて一枚も捨てないようなもので、机がどれほど大きくても足りなくなります。
そのため、計算を減らすだけでなく、保存も減らすことが必要です。
文心Unlimited OCR 2はこのために無限デコード機構(Unlimited-Horizon Decoding Paradigm)を導入し、チャンク処理によってキャッシュ占有を制御します。

△チャンクデコードの模式図
この機構は文書をサイズが制御可能なページブロックに分割し、順次解析します。隣接する2つのブロックは1ページ重なり、前のブロックの最終ページは次のブロックの起点にもなります。まさにそのまま使えるしおりです。
次のブロックに入ると、モデルはこの重複ページの視覚KVキャッシュを直接再利用するとともに、前のセグメント出力末尾のテキストキャッシュも保持し、前後の内容がつながるようにします。それ以外のもう使わないページの視覚キャッシュは速やかに解放し、VRAMを確保します。
こうして、VRAMに保持すべき視覚キャッシュは、現在処理中のブロックの規模内に抑えられます。1ブロックの処理が終われば、接続に必要な情報だけを残して次のブロックを開始します。文書全体の視覚キャッシュを最初から最後まで占有させる必要はありません。
*SR-SWAは各ステップの計算量を削減し、無限デコード機構はキャッシュ占有を制御します。両者が協働することで、モデルは超長文書を継続的に解析できます。884ページの『ロード・オブ・ザ・リング』を42分で読了、キャッシュは1.77GBに低減 技術についてはこれだけ説明してきましたが、私たちもその効果をこの目で確かめたいところです。早速、884ページの英語版『ロード・オブ・ザ・リング』。

を渡してみたところ、約42分56秒後に解析が完了しました。884ページの入力に対して884ページの出力で、1ページあたり平均約2.9秒です。テレビドラマ1話分ほどの時間で、「中つ国」の旅が一周し終えたことになります。 これほど長い文書でも、後半になっても識別精度は維持できるのでしょうか? 私たちは第600ページと第750ページの本文を特別に抜き取り確認しました:識別テキストは原文と一致しています。第750ページに至っても一字一句違わないということは、モデルが前半で「読みすぎて」いないことを示しています。

このタスク全体を通じて、システムはさらに466カ所のページ跨ぎ結合をマークしました。たとえば本の終わりに近い第873〜874ページでは、ページ分割によって無理やり切断された文も、出力の中で再びつながりました。

実測では本全体の連続解析の過程を示し、技術レポートはさらに3つの問いに答えています:どれだけのキャッシュが必要か、どれだけ速くできるか、識別品質はどうか。百ページ規模のベンチマークMarathonDocBenchにおいて、チームは標準デコードとUnlimited-Horizon Decodingを並べて比較しました。単一系列推論では、チャンク処理と境界キャッシュ再利用を採用した結果、ピーク時のKVキャッシュ占有が47GBから1.77GBに低下し、約96%削減されました。*
従来はVRAMに同時に残しておく必要があった大量のページ情報を、現在は分割処理して速やかに解放できます。ここの1.77GBはKVキャッシュのみを指し、モデルの重みなどの他のVRAM消費は別途計上されます。
キャッシュの節約に加えて、レポートでは複数文書の同時処理時の性能も測定しました。
別の16路同時実行テストでは、このデコード方式を採用後、エンドツーエンドのスループットが360から1315Token/秒に向上し、標準デコードの約3.65倍となりました。

△デコード方式の比較:単一系列のピークKVキャッシュと16路同時実行のスループット
年次報告書や技術マニュアルを大量に処理する必要があるサービスにとって、この2つの変化はいずれも実用的な意味を持ちます:VRAM占有が下がれば、同じ1枚のカードでより多くの仕事をこなせますし、同時実行が走れば、単位時間あたりの出力も増えます。
ユーザーは何も気にする必要がなく、従来通り文書全体を提出すれば、チャンク分割と接続はモデルが内部で行い、手動でのファイル分割、分割提出、再結合の手間が省けます。
より速く、より省エネであることに加えて、解析品質も守られなければなりません。
MPDocBenchは複数ページ文書の解析品質を評価します:数十ページに及ぶ文書に対して、モデルは内容を正確に識別すると同時に、文書構造も再現しなければなりません。
このテストにおいて、文心Unlimited OCR 2は84.59の総合スコアを獲得しました。
百ページ規模の文書をさらにテストするため、チームは新たにMarathonDocBenchを作成した。これには100件の文書が含まれ、それぞれ103~300ページである。多くの段落や表が次のページにまたがって続き、モデルが前後の内容を正しくつなげられるかを重点的に試すものである。
このより長く、より手強いベンチマークにおいて、文心Unlimited OCR 2は77.06の総合スコアを達成した。両方の評価において、その総合スコアは報告書に記載された参戦モデルの中で最高値となった。

△技術報告より:複数ページ文書解析の評価結果
長文書処理速度の変化も同様に顕著である。報告書の百ページ規模テストにおいて、SR-SWAの1ステップあたりのデコード速度は、Full Attention(フルアテンション機構)の7.0倍、前世代で使われたR-SWAの5.6倍。
に達した。グラフを見るとより直感的である。文書が長くなるにつれて、前者2つの方式の1ステップデコード時間は上昇し続けるが、SR-SWAの1ステップ時間はほぼ水平線である。
これは、モデルが1つのTokenを生成するのに必要な時間は、入力ページ数の増加によって明確に長くならないことを意味する。本が厚くなっても、書く速度は変わらないのである。

△異なる文書長における1ステップデコードのレイテンシ
前世代と比較すると、文心Unlimited OCR 2の1ページあたりの総合スコアはほぼ同等である。
本当の差が開くのは、文書が長くなってからである。連続処理できるページ数がより多く、ページをまたぐ内容の認識がより正確になり、長文書解析の計算コストとキャッシュコストも同時に低下する。
皆さんがどう使うかというと……
数百ページの年次報告書や、分厚い技術マニュアルを入力できる。
特にスキャン文書の場合:まずモデルに連続解析させて文字に変換し、その後情報の検索や内容の照合を行えば、ページの分解や手作業での継ぎ合わせの手間も大幅に減らせる。
しかし、これほど強力なOCRを、仕事で使う人がいるだろうか……

もちろん、小説を読むためです!

OCR 3.0時代、到来
もちろん、数ページの資料から書籍一冊まで、OCRが解決すべき課題も変化している。
技術報告の中で、チームは文書解析の進化をOCR 1.0、2.0、3.0の3段階に概括した。
初期のOCRは、役割分担がはっきりした一つの流れ作業のようなものだった。まず文字の位置を見つけ、次にその内容を認識し、最後にレイアウトを整理して結果を出力する。複数ページの解析が絡む場合は、通常1ページずつ処理し、その結果をつなぎ合わせていた。これがOCR 1.0である。
各モジュールはそれぞれの役割を担っているが、前の段階での誤りがそのまま後ろへ伝わることもある——一つの工程がくしゃみをすれば、後ろはすべて風邪をひくのだ。
到了OCR 2.0では、視覚言語モデルがページ画像から直接解析結果を生成するようになり、複数モジュール間の受け渡しを減らした。初代のUnlimited OCRは、このエンドツーエンドの解析を数十ページの文書へと拡張した。

△AI生成
文書が数十ページから数百ページへと増える場合、連続解析では前後の内容のつながりを保ちながら、計算とキャッシュのコストも抑える必要がある。
文心Unlimited OCR 2はオンデマンド参照とチャンク単位のデコードにより、より多くのページを処理する際に、各ステップの計算範囲がページ数に応じて拡大することも、すべてのページの視覚キャッシュを保持し続けることも不要にした。
チームはこのステップを次のように呼ぶOCR 3.0:文書全体を継続的に処理することを新たな技術目標とする——一気に、ページ分割なしで、無限に解析し、決してダウンしない。
しかも、この目標には二次元バージョンまである。
Unlimited OCR Worksという名前自体、《Fate/stay night》の「無限剣製」、すなわちUnlimited Blade Works。

△《Fate/stay night》アニメ公式サイトトップページ
原作の設定では、これは「固有結界」と呼ばれる魔術で、発動すると無数の刀剣が突き刺さった特殊な空間が展開される。
前世代のリリース後、Hugging FaceユーザーのSnarcyがこの命名ネタに気づき、定番の詠唱をOCR版に書き改めた。
第二代では、チームは実際にこの書き換えを技術レポートの序文に収め、Snarcyの署名も残した。

こうして、アテンション機構とVRAMコストを論じるこのレポートを開くと、冒頭で次のような詠唱を目にすることになる:
ピクセルは私の身体であり、トークンは私の血である。
body(身体)はピクセルにして、血はToken。
I have parsed over a thousand pages.
我はすでに千余ページの文書を読み尽くした。
Unknown to KV cache.
KV Cacheに知られず、
Nor known to context limits.
コンテキスト長の制限にも縛られぬ。
So, as I pray: Unlimited OCR Works
ゆえに、我はここに祈る:Unlimited OCR Works。

47GBから1.77GBへ、「KVキャッシュに縛られぬ」という一句、唱えた甲斐があったというものだ。
これもすっかり中二病である。

—完—

@量子位 · AI技術と製品の最新動向を追いかける

共感していただけた方は、高評価・フォロー・シェアの三連を歓迎しますվ’ᴗ’ ի ❤

原文の出典

量子位

内容について

原文の公開と権利は出典元に帰属します。

機械翻訳 · 原文をご参照ください