はじめに

こんにちは。機械学習エンジニアの上川です。LayerXでは、バクラクのAI-OCR機能の精度改善に取り組んでいます。本記事では、Data-centric AIにまつわる技術を用いて、AI-OCRデータセットの品質改善を行うための技術検証を行なったのでその紹介をします。

AI-OCRデータセットのアノテーション

AI-OCRとは、国税関係書類に対してOCRを実行し、入力のサジェストを行うことでユーザーが書類の内容を手入力する手間を削減する機能です。例えば、領収書の日付、金額、取引先名を自動で読み取り、ユーザーにサジェストします。

AI-OCRでは内製の機械学習モデルを利用しているため、学習用・検証用のデータセットにアノテーションを行い、AI-OCRデータセットを作成する必要があります。ここでは、画像内のテキストを抽出するだけでなく、項目の意味や位置の情報をテキストに付与しています。例えば、請求書の場合、テキストに対して請求日や請求金額などの意味や位置をアノテーションします。

アノテーションを一部自動化する工夫もしていますが、項目の意味や位置という2種類の情報をテキストに付与する必要があるので、単純な分類タスクなどと比較するとアノテーションの負荷は高くなってしまいます。

tech.layerx.co.jp

そのため手作業でアノテーションを行う場合、誤りが生じる可能性があったり、テキストの意味の解釈が曖昧な場合などでノイズが混入することもあります。

したがってアノテーションの負荷や難しさなどの観点から、AI-OCRデータセットの品質向上は、AI-OCRの精度改善において重要な課題となっています。

Data-centric AI

データセットの品質向上において、近年はData-centric AIというアプローチが注目されています。 Data-centric AIとは、従来のModel-centric AIとは異なり、モデルの改善ではなく、データの品質改善に焦点を当てたAI開発のアプローチです。

www.youtube.com

Data-centric AIの単純なアプローチを考えると、データを目視することによって、ノイズを抽出し改善していく方法が挙げられますが、数千枚~数万枚のデータを人手でチェックするのは非常にコストがかかります。また、人手による確認では、正解が曖昧な場合に主観的な判断が入ってしまいます。

これらの課題を改善するためにConfident Learningという手法が提案されています。Confident Learningは、モデルの予測結果とアノテーションの間に矛盾があるデータをノイズとして検出し除去することで、データセットの品質を向上させる手法です。

Cleanlabを用いたノイズ検出

Confident Learningを利用したノイズ検出のためのツールとして、Cleanlabというライブラリが提供されています。*1

AI-OCRでは、冒頭でも述べたようにテキストの意味と位置の2種類の情報をアノテーションをする必要があるのですが、本記事ではテキストの意味のアノテーションに着目し、Cleanlabを用いてノイズを検出します。

Cleanlabを用いたノイズ検出は様々な種類のデータセットに適用可能ですが、固有表現抽出（Named Entity Recognition, NER）タスクにConfident Learningを適用することで、テキストの意味におけるノイズ検出を行います。*2

検証用の機械学習モデルとしてtransformersのRoBERTa*3、検証用のデータセットとしてCoNLL-2003*4を利用します。

以下、検証の手順です。

1. データセット準備

CoNLL-2003データセットをダウンロードし、学習データ、検証データ、テストデータに分割します。

from datasets import load_dataset

dataset = load_dataset("conll2003")
train_data = dataset['train']
val_data = dataset['validation']
test_data = dataset['test']

label_list = dataset['train'].features['ner_tags'].feature.names

CoNLL-2003では、テキスト上の固有表現に対して、persons(人物)、organizations(組織)、 locations(場所)、 miscellaneous(その他)のラベルがついています。また固有表現抽出のラベル付け方法はいくつかありますが、BIO形式という方法が用いられています。この方法では、固有表現の始端のトークンにB-(固有表現の種類)、それ以降の固有表現内トークンにI-(固有表現の種類)、固有表現以外のトークンにはOのラベルが付与されます。

Yuta	Kamikawa	works	as	a	ML	engineer	for	LayerX	.
B-PER	I-PER	O	O	O	O	O	O	B-ORG	O

2. ファインチューニング

transformersのRoBERTaをNERタスク用にファインチューニングします。

import os
from transformers import RobertaForTokenClassification, RobertaTokenizerFast, Trainer, TrainingArguments, DataCollatorForTokenClassification

model_name = "roberta-base"
model = RobertaForTokenClassification.from_pretrained(model_name, num_labels=len(label_list))
tokenizer = RobertaTokenizerFast.from_pretrained(model_name, add_prefix_space=True)

# Tokenization and alignment of labels function
def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(examples['tokens'], truncation=True, is_split_into_words=True, padding=True)
    labels = []
    for i, label in enumerate(examples[f'ner_tags']):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx is None:
                label_ids.append(-100)
            elif word_idx != previous_word_idx:
                label_ids.append(label[word_idx])
            else:
                label_ids.append(-100)
            previous_word_idx = word_idx
        labels.append(label_ids)
    tokenized_inputs["labels"] = labels
    return tokenized_inputs

# Tokenize and align labels for train and validation, test datasets
tokenized_train_data = train_data.map(tokenize_and_align_labels, batched=True)
tokenized_val_data = val_data.map(tokenize_and_align_labels, batched=True)
tokenized_test_data = test_data.map(tokenize_and_align_labels, batched=True)

# Use DataCollatorForTokenClassification to handle padding of inputs and labels
data_collator = DataCollatorForTokenClassification(tokenizer)

# Training arguments
training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
)

# Trainer object
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train_data,
    eval_dataset=tokenized_val_data,
    data_collator=data_collator,
)

# Start training
trainer.train()

3. 予測確率の取得

ファインチューニングしたRoBERTaを用いて、CoNLL-2003の学習データに対する予測確率を取得します。予測確率は、各トークンが各ラベル(人物、場所、組織など)に属する確率を表すベクトルとなっています。

# Trainerのpredictメソッドを使って予測とラベルを取得
def get_predictions_and_labels(trainer, dataset):
    predictions_output = trainer.predict(dataset)
    
    # logits (予測) と labels を取得
    logits = predictions_output.predictions
    labels = predictions_output.label_ids
    
    # logitsにSoftmaxを適用して確率を取得
    probabilities = torch.softmax(torch.tensor(logits), dim=-1).numpy()
    
    return probabilities, labels

# train, validation, testデータに対して予測とラベルを取得
train_predictions, train_labels = get_predictions_and_labels(
    trainer,
    tokenized_train_data
)

4. ノイズ検出

取得した予測確率とCoNLL-2003データセットのアノテーション（正解ラベル）をCleanlabに入力し、ノイズ検出を行います。find_label_issues()を使用して、ノイズの可能性が高いトークンを特定することができます。この時予測確率と正解ラベルとの矛盾が大きいデータをノイズとして抽出します。

import numpy as np

from cleanlab.token_classification.filter import find_label_issues
from cleanlab.token_classification.rank import get_label_quality_scores

labels = []
pred_probs = []
pred_labels = []

for label_seq, pred_seq in zip(train_labels, train_predictions):
    valid_labels = []
    valid_predictions = []

    for label, pred in zip(label_seq, pred_seq):
        if label != -100:
            valid_labels.append(label)
            valid_predictions.append(pred)
    
    labels.append(np.array(valid_labels))
    valid_predictions = np.array(valid_predictions)
    pred_probs.append(valid_predictions)
    pred_labels.append(np.argmax(valid_predictions, axis=-1))

# Cleanlabを使ってノイズを検出
label_issues = find_label_issues(
    labels=labels,
    pred_probs=pred_probs,
    return_indices_ranked_by="self_confidence"
)

# 結果の出力
print("Label Issues:", label_issues)

5. ノイズ分析

display_issues()とcommon_label_issues()を使用することによって、検出されたラベルの分析を行うことができます。

from cleanlab.token_classification.summary import display_issues, common_label_issues

display_issues(
    label_issues,
    train_data["tokens"],
    labels=labels,
    pred_probs=pred_probs,
    exclude=[(0, 1), (1, 0)],
    class_names=label_list
)

info = common_label_issues(
    label_issues,
    train_data["tokens"],
    labels=labels,
    pred_probs=pred_probs,
    exclude=[(0, 1), (1, 0)],
    class_names=label_list,
)

こちらがノイズとして検出されたデータの例です。

CoNLL-2003データセットは、ベンチマークとしてよく利用されるデータセットですが、これらのようなノイズが含まれていることがわかりました。また検出されたノイズを分析すると、人名を地名やその他のラベルと間違えているパターンが多くみられます。固有表現抽出のアノテーションには、人名や地名などの前提知識が必要なことから、これらのノイズが発生したと考えられます。

6. 再学習、評価

ノイズとして検出されたデータを除去し、再学習と評価を行います。

# 除去後のデータセットを用いて再度ファインチューニングを行います。
trainer.train()
trainer.evaluate(eval_dataset=tokenized_test_data)

こちらがノイズ除去前後の性能比較になります。僅かではありますが、Confident Learningによりノイズと判定されたデータを取り除くことで、すべての項目において性能が向上するという結果になりました。

Category	Metric	除去前	除去後	Difference
LOC	precision	0.936043	0.939157	0.00311359
LOC	recall	0.938849	0.944652	0.00580336
LOC	f1	0.937444	0.938899	0.00145516
MISC	precision	0.782723	0.799737	0.0170143
MISC	recall	0.851852	0.863333	0.0114815
MISC	f1	0.815825	0.860274	0.0444482
ORG	precision	0.893567	0.896244	0.00267665
ORG	recall	0.919928	0.922938	0.00301023
ORG	f1	0.906556	0.907296	0.000740269
PER	precision	0.973091	0.973225	0.000134041
PER	recall	0.961657	0.966605	0.00494743
PER	f1	0.967341	0.969904	0.00256323
overall	precision	0.913316	0.921621	0.00830551
overall	recall	0.929001	0.937762	0.00876062
overall	f1	0.921092	0.939621	0.018529
overall	accuracy	0.983676	0.984116	0.000440078