バクラクOCRにおける機械学習モデルの検証から本番投入までの取り組み

機械学習エンジニアの吉田です。今年の3月にLayerXに入社し主に機械学習を活用したOCRの開発をしています。それまでルールベースで実装されていた工程を一部機械学習ベースの実装に置き換えつつあり実際に本番でも稼働しています。この記事ではどのように検証や開発を進めてきたのかご紹介します。

背景

LayerXで提供しているバクラクでは帳票をアップロードするだけで支払金額や支払期日などを自動で読み取り補完してくれるOCR機能があります。このOCRは主にルールベースで実装されていましたが、リリースから継続的に改善が行われていたことでかなり精度高く読み取ることができるようになっていました。しかし導入企業様が増えるにつれて次第に精度の改善が困難になり、一箇所修正するとこれまで読み取れていたものが読み取れなくなるなどコードの修正も困難になってきている状況でした。一方でリリースから1年以上経過し集められた帳票の数もかなり多くなってきていたので機械学習ベースのOCRを作ることでこれらの問題を解決できるのではないかという期待がありました。その辺りの背景についてはこちらの記事に詳しく書かれているのでぜひご覧ください。 note.com note.com

OCRは２つの処理に大別されます。

帳票に書かれている文字列を認識し検出すること
検出された文字列から必要な項目を抽出すること

今回機械学習によって解きたい問題は２つ目の検出された文字列から支払金額や支払期日などの項目を抽出するところになります。

データセットの整備

元々社内には帳票に対して支払金額や支払期日など人手でアノテーションされたデータがあったので最初はこのデータを使ってモデルの学習を行う想定でいました。しかし、当時はアノテーションデータにはbounding boxがなかったり、学習データとして使うにはデータ件数も多くはありませんでした。(ちなみに現在はbounding boxもアノテーション可能な管理画面v2が稼働しています)

まずはモデルの精度よりもルールベースやモデルの正確な評価を優先したかったのでアノテーションデータは全て精度評価の為のテストデータとして確保し、学習データにはルールベースの結果やユーザの入力値を使うことにしました。項目によってはルールベースの精度はかなり良かったのでルールベースの結果を正解として扱ってもある程度モデルを学習させることができるのでは、という期待がありました。過去の文字検出やルールベースによる項目推定の結果はすべてJSONファイルで保存されていたのでこれをパースすることで大量の学習データセットをすぐに作成することができました。

また、アノテーションデータからはテストデータを作成しルールベースにおける各項目の精度(Accuracy)を出せるようにしました。この精度が今回の機械学習モデルの目標とすべき精度となりました。

ベースラインモデルの構築

機械学習の開発は不確実性が大きいのでまずはパイプラインを一気に作ってしまい見込みがあるのかどうか、本番運用に問題がないかを早い段階で判断し、改善のイテレーションを高速で回すことが重要です。まずはあまり深く考えずにベースラインのモデルを作ってルールベースの精度と比較ができるところまでもっていくことを優先しました。

検出された文字列から項目を抽出するにあたり機械学習に解かせるタスクとしてはいくつか考えられますが、まずはシンプルにトークン分類を試してみました。最初のモデルとしてはLayoutXLMを採用しました。LayoutXLMは請求書等の大量の文書画像で事前学習されたモデルで、多言語の事前学習済みモデルも公開(商用利用は不可)されていたので初手に試すのに最適でした。

トークン分類の出力はそのままではOCRの結果としては使うことができません。どのトークンをひとつづきの文字列として扱うのか、同じ項目が複数箇所で検出された場合にどれを採用するか等の後処理の工夫が必要となります。ひとまずはその辺の後処理も後で考えることにしてナイーブな実装としました。ここまでで最低限ルールベースの出力と揃えることができたので、精度を評価してみたところルールベースには遠く及ばないもののちゃんと学習できていることが分かりました。

続いて推論結果をエラー分析してみたところ推論された箇所はあっているのに微妙に正解と異なるものが大量にあることがわかりました。例えば、推論結果の支払金額にカンマが含まれてしまっていたり、支払期日が「翌月末」となっているが「2022/10/31」と日付を期待しているケースです。モデルでは対応が困難なこれらのケースを後処理に追加してみたところ大幅に精度を改善することができたので方向性は間違っていないという確信をもつことができました。

学習から推論、評価まで一通りのパイプラインの構築はできたので、次に実際の本番運用を見据えて主に推論時の問題を解決することを優先しました。もし何かクリティカルな問題があった場合、モデルを改善したとしてもリリースすることができず手戻りが大きくなってしまうので早めに対応しておく必要があります。

例えば文字数が多い帳票の場合、一度にモデルに入力できる長さに制限があるため、複数回推論してマージするといった処理が必要となりますが、その場合に精度は問題がないのか？本番でのオンライン推論時の実行時間に問題はないか？といったような問題です。

上記のような推論時の問題を解消するのと並行して学習データ作成の見直しも行いました。学習データはルールベースの結果を元に生成していたこともあり、どうしてもノイズが入ってしまっていたり、そもそも全然違う文字列が正解となっていたりしました。学習データの件数を減らしてでもノイズが少なく精度の高いデータセットを用意した方が精度を改善できることが分かりました。

ここまで主に後処理と学習データの改善だけで、ルールベースの精度にだいぶ近づけることができましたが、ほぼ改善の手を打ち尽くしたのでいよいよモデルの改善に着手することにしました。

実験管理

まず最初に行ったのは、実験を管理することができるようにすることと再現性の担保と精度のばらつきを確認することでした。実験管理にはWandBを採用しました。前職ではMLflowを自前で立てていましたが、WandBではインフラが不要ですぐに使い始めることができるので少人数で使うには便利です。

また、seedや学習データを様々に組み替えて実験してみたところ、思った以上に精度にばらつきが発生することが分かりました。 OptimizerやSchedulerを調整することで学習の安定化と学習時間を短縮し、学習データとテストデータの分割を工夫してロバストに評価できるようにするなど実験の土台を整えることで今後のモデルの改善を効率的に行えるようにしました。

モデル改善

ベースラインとして利用していたLayoutXLMの事前学習モデルは商用利用することができないので、本番に投入することのできるモデルを２つの方向性で検討することにしました。

ひとつはLayoutLMの事前学習モデルを自前で作ることです。それについてはこちらの記事に詳しく書かれていますのでご覧ください。 tech.layerx.co.jp

もうひとつはRoBERTaの事前学習モデルをベースとすることです。RoBERTaはLayoutLMのようにbounding boxや画像の特徴量を使わないので精度の面で懸念はありましたが、これまでに作ったパイプラインをそのまま使ってすぐに改善を回し始めることができるというメリットがありました。

ここまでモデルには手を入れずに、学習データや後処理の改善に注力していたことが功を奏しRoBERTaでも最初から良い精度を出すことができました。さらにRoBERTaの精度を改善する上で効果の大きかったものには以下のようなものがありました。

Label Smoothingによる正則化
- 前述したように学習データにはノイズが含まれていたので過学習抑制の為にLabel Smoothingを入れてみたところ精度が改善されました
損失関数でクラスの重みを調整
- 支払金額や支払期日のような目的とするクラス以外はすべて「その他」として扱っていたので、当然「その他」に属するトークンが圧倒的に多い状況でした
- 損失関数において「その他」の重みを調整することで各項目のprecisionとrecallを動かすことができたので、後処理の対応がしやすい項目はrecallを高めることで前後の余分なトークンが含まれていたとしても精度を改善することができました
task-adaptive pretraining (https://arxiv.org/abs/2004.10964)
- 事前学習モデルに対してバクラクのデータセットで追加で事前学習をしたところ精度を大きく改善することができました
- 追加で事前学習させるデータを増やすほどファインチューニング後の精度は上がり、約20万件で飽和しました。
- 20万件の事前学習にかかった時間はGCPでA100 x 16インスタンスを使って20時間程度でした