LayerX エンジニアブログ

LayerX の エンジニアブログです。

LLMを活用した機械学習モデルのアノテーション効率化

機械学習 LLM

機械学習エンジニアの吉田です。今回は、LLM (Large Language Models) を活用して、機械学習モデルに必要なデータのアノテーション作業を効率化する取り組みについて紹介します。

なお、アノテーションにおけるLLMの利用に関しては、クラウドベンダー各社及び社内の法務確認のうえ進めています。この記事で登場するLLMがアノテーション用途で使えることを保証するわけではないのでご留意ください。

背景

LayerXで提供しているバクラクでは、帳票をアップロードするだけで支払金額や支払期日などを自動で読み取るOCRの機械学習モデルを開発しています。

このモデルのデータセットを作成するために、内製のアノテーション基盤を使って日々手作業でアノテーションを行っています。帳票の種類が多く、解釈が複数ある場合もあるため、アノテーション作業は簡単ではありません。モデルの推論結果やユーザーの入力値を候補としてサジェストしたり、バウンディングボックスで囲った領域の文字を自動入力するなど、アノテーターの作業負荷を軽減する工夫をしています。

tech.layerx.co.jp

既存の項目をアノテーションするだけであれば効率的に作業ができる仕組みは整えられていましたが、今後アノテーションする項目を大幅に増やしたいと考えており、そのときに以下の課題がありました。

1. アノテーション作業時間の増加

  • アノテーション項目が増えることで作業時間が増加
  • 新規項目に対して推論可能なモデルやユーザーの入力値がなく、候補をサジェストできないため難易度が上がる

2. 機械学習モデルの検証に時間がかかる

  • アノテーションに時間がかかることで、機械学習モデルの学習・評価に必要なデータセットの作成に時間がかかる
  • モデルの性能がでなかった場合にアノテーションの見直しが必要になると手戻りが大きくなる

3. システムの拡張性が無い

  • アノテーション項目の変更に対してDBやアプリケーションが柔軟に対応できない

LLMによるアノテーション

LLMは様々な自然言語処理タスクで優れたゼロショット能力を示しています。大量の未ラベルデータに対してLLMでアノテーションし、BERTのような小規模なモデルを学習させる研究も多くみられます。*1 *2

LLMを活用することで、ゼロショットでアノテーション候補をサジェストしたり、LLMの推論結果を疑似ラベルとしてモデルをトレーニングすることで、早い段階でアノテーションの方向性に問題がないか検証できるのではないか、という期待があります。

一方でLLMを活用する際には、構造化データをロバストに抽出する方法や、システムの拡張性の課題に加えて、プロンプトを柔軟に変更できる仕組みが必要です。

JSON Schemaによるスキーマ管理

現行のアノテーション基盤では、各項目に対してDBのカラムを設けているため、項目を追加する際にはDBスキーマの更新が必要でした。

これに対し、新しいアノテーションスキーマでは階層構造を持たせ、構造の変化にも柔軟に対応できるよう、DBには正規化を行わずにJSONで保存する方法を採用しました。これにより、柔軟なデータ構造に対応できるようになりましたが、その一方でデータの整合性を保つためには、スキーマの管理やデータのバリデーションが重要になります。

バクラクではスキーマ管理にProtobufを使用していますが*3、今回はJSON Schemaを採用しました。 JSON Schemaを選んだ理由は、Amazon BedrockやAzure OpenAIなどの主要なAPIでリクエストパラメータとしてJSON Schemaを指定できるためです。JSON Schemaを使ってLLMから情報抽出できれば、プロンプトとDBに保存するJSONのスキーマを一元管理できると考えました。

JSON SchemaによるLLM推論

Amazon BedrockやAzure OpenAIではFunction callingでJSON Schemaを使うことができます。 Amazon Bedrockの場合はConverse APIの tools パラメータで*4、Azure OpenAIの場合はChat Completions APIの tools パラメータで*5 JSON Schemaを指定できます。

先日OpenAIからはJSON Schemaの構造で抽出するための Structured Outputs*6 がリリースされましたが、JSON Schemaの一部仕様に対応していなかったり、ネストの階層に制限があったり、予期しないエラーが発生することがあり、私たちのユースケースには適していませんでした。

以下は簡略化したJSON Schemaの例です。書類の発行日 issueDate と支払金額 paymentAmount の2つのプロパティがあり、それぞれに、rawText と processedValue というプロパティを持っています。 rawText には書類に記載された文字列が入り、processedValue にはISOフォーマットの日付や数値に変換された金額が入ります。

LLMには、情報の抽出に専念させたいことと、ハルシネーションを避けるため、processedValue の後処理はできるだけ避け、型やフォーマットの変換にとどめます。そして、processedValue の値はバリデーションでチェックします。
また、additionalProperties や required プロパティについても、LLMでは基本的に遵守できませんが、これらもバリデーションでチェックします。

また、LLMに渡す前に、$ref 参照を解決し、空白や改行を取り除いたJSON Schemaに変換しています。その際、アノテーション基盤だけで使用する管理用のプロパティ metadata や、LLMでは推論できない位置情報 position も削除しています。

{
  "$schema": "http://json-schema.org/draft-07/schema#",
  "additionalProperties": false,
  "properties": {
    "issueDate": {
      "additionalProperties": false,
      "description": "書類の発行日",
      "properties": {
        "processedValue": {
          "description": "書類に記載されている日付をISO 8601フォーマットに変換した日付",
          "format": "date",
          "type": [
            "string",
            "null"
          ]
        },
        "rawText": {
          "description": "書類に記載されている日付",
          "type": [
            "string",
            "null"
          ]
        },
        "position": {...},
        "metadata": {...}
      },
      "required": [
        "rawText",
        "processedValue",
        "position",
        "metadata"
      ],
      "type": "object"
    },
    "paymentAmount": {
      "additionalProperties": false,
      "description": "支払金額",
      "properties": {
        "processedValue": {
          "description": "書類に記載されている金額を数値に変換",
          "type": [
            "number",
            "null"
          ]
        },
        "rawText": {
          "description": "書類に記載されている金額",
          "type": [
            "string",
            "null"
          ]
        },
        "position": {...},
        "metadata": {...}
      },
      "required": [
        "rawText",
        "processedValue",
        "position",
        "metadata"
      ],
      "type": "object"
    }
  },
  "required": [
    "issueDate",
    "paymentAmount"
  ],
  "type": "object"
}

バリデーション

アノテーション基盤とLLM推論のバッチ処理はGoで開発しているため、JSON Schemaを用いたバリデーションには gojsonschema を、JSON SchemaからGoの構造体を生成するために schematyper を使いました。

バリデーションが必要な場面は2つあります。1つは、人間がアノテーション作業を行うとき、もう1つはLLMの推論結果を保存するときです。アノテーション作業では、JSON Schemaに基づいたバリデーションを行えばよいですが、LLMの結果を保存する際には少し工夫が必要です。LLMの推論結果でバリデーションエラーが発生する可能性があるケースは以下の通りです。

  1. プロパティの型が違う
  2. JSONが無効
  3. 定義にないプロパティが含まれている
  4. 必須のプロパティが欠けている

1〜3のケースについては、リトライすることでバリデーションをパスすることができますが、4については必須とされているプロパティがどうしても欠けてしまう場合があります。
そこで、LLMの推論結果をバリデーションする際には、まず必須項目以外をチェックし、問題がなければ、欠けているプロパティを補完して初期値を設定します。その後、再度必須項目を含めたバリデーションを行うようにしています。

バージョン管理

アノテーションのスキーマは、将来的に変更される可能性があります。もし、単純にJSON Schemaのプロパティが追加・削除されるだけであれば、古いバージョンのJSON Schemaで作成されたJSONは、新しいバージョンのJSON Schema（具体的には、JSON Schemaから生成されたGoの構造体）で読み込むことが可能です。ただし、プロパティの型が変更されたり、スキーマ内で位置が変わったりして、以前のアノテーションの値を保持したい場合には、バージョン管理が必要となります。

バージョン管理が必要な場合でも、JSON Schema (schema.json) を以下のようにバージョン毎のディレクトリに分けるだけで対応できるようにしています。schema.gen.json と schema.gen.go は、JSON Schemaから自動生成されます。 また、JSONを構造体に変換する json パッケージや、バリデーションを行う validator パッケージは、バージョンに依存せずに利用できるようになっています。

.
├── json
│   └── json.go  // JSONを構造体に変換
├── schema
│   ├── v1
│   │   ├── schema.gen.go    // 自動生成した構造体
│   │   ├── schema.gen.json  // 参照解決と空白・改行、不要なプロパティを削除、プロンプト用
│   │   └── schema.json
│   └── v2
│       ├── schema.gen.go
│       ├── schema.gen.json
│       └── schema.json
└── validator
    └── validator.go  // バリデーション

// json.go
func Unmarshal[T any](data []byte) (*T, error) {
    var a T
    dec := json.NewDecoder(strings.NewReader(string(data)))
    err := dec.Decode(&a)
    if err != nil {
        return nil, fmt.Errorf("failed to unmarshal JSON: %w", err)
    }
    // 初期値のセット
    setDefaults(&a)
    return &a, nil
}

// validator.go
func Validate(schemaData []byte, inputData []byte) *ValidationError {
    schemaDataLoader := gojsonschema.NewBytesLoader(schemaData)
    inputDataLoader := gojsonschema.NewBytesLoader(inputData)
    result, err := gojsonschema.Validate(schemaDataLoader, inputDataLoader)
    // 省略
}

まとめ

本記事では、アノテーション作業を効率化するために、どのようにLLMを取り入れているかを紹介しました。

JSON Schemaを用いることで、LLMプロンプトとDBに保存するJSONを統一し、変化するアノテーションのニーズに柔軟に対応できる仕組みを構築しています。
また、バリデーションのロジックを統一することで、データの整合性を保ちつつ、安全にアノテーション作業を進めることが可能となっています。

今後の課題としてはLLMプロンプトのチューニングです。スキーマ定義がLLMの性能にどのような影響を与えているかを定量的に評価し、最適な構造を模索しながら改善を続けていく予定です。

*1:https://arxiv.org/abs/2403.15938

*2:https://arxiv.org/abs/2402.15343

*3:https://tech.layerx.co.jp/entry/2023/11/16/140041

*4:https://docs.aws.amazon.com/bedrock/latest/userguide/tool-use.html

*5:https://learn.microsoft.com/ja-jp/azure/ai-services/openai/how-to/function-calling

*6:https://openai.com/index/introducing-structured-outputs-in-the-api/