マルチモーダルな商品カテゴリの分類モデル

はじめに

この記事は、BASEテックブログ夏のブログリレー10日目の記事です。

こんにちは、Data Strategyチームの竹内です。

BASEでは日々数多くの多種多様な商品が新しく登録されています。それらの商品が「何のカテゴリの商品なのか」を機械学習モデルで自動的に推論する仕組みを、以前からバッチ処理基盤として運用してきました。

今回、そのモデルを 商品テキストと商品画像の両方を入力に取るマルチモーダルなモデル に置き換えたので、その経緯やモデルの中身について紹介します。

はじめに
TL;DR
なぜ商品カテゴリを推論するのか
これまでの取り組みと、画像を使いたくなった理由
使用したモデル
結果
埋め込みベクトルの活用
おわりに

※ 記事内のコードはサンプルとして簡略化しています。

TL;DR

商品タイトル・説明文などのテキストと、商品画像の両方を入力とする商品カテゴリ分類モデルを作成しました
モデルは MMBT（Supervised Multimodal Bitransformers）をベースに、画像側を Swin Transformer、テキスト側を日本語BERT（tohoku-nlp/bert-base-japanese-v3）に置き換えたものです
学習データは約100万件の商品で、ラベル（498クラスの階層カテゴリ）はLLMによるアノテーションで付与しました

なぜ商品カテゴリを推論するのか

BASEには毎日さまざまな商品が登録されますが、「その商品がどのカテゴリに属するか」を横断的に把握することは困難です。ショップ側で設定されたショップカテゴリは任意項目であり、また「ショップのカテゴリ」と「そのショップが実際に売っている個々の商品のカテゴリ」は必ずしも一致しません。

そこで、商品ごとにカテゴリを機械学習で推論して付与しています。用途は大きく3つあります。

分析基盤としての提供 … カテゴリごとの流通額・登録数・不正決済の発生状況といった、商品全体像に対する解像度を上げるための分析軸
不正検知モデルの特徴量 … 不正決済の検知や不正な商品登録の検知など、各種モデルの特徴量として利用
プロダクトでの活用 … 推薦や検索など、アプリ側の機能での利用

これまでの取り組みと、画像を使いたくなった理由

商品カテゴリの推論そのものは新しい取り組みではなく、2022年にBERTを使ったモデルとその推論基盤について記事を書いています。

devblog.thebase.in

このときのモデルは商品タイトルと説明文を結合したテキストのみを入力とするもので、記事でも画像の利用は今後の課題としていました。

商品によっては、説明が簡素で短いものや、購入者とのやり取りや注意事項のみを記載しているものなど、テキストだけでは何の商品か判断できない商品が一定数あります。一方で、そうした商品でも画像を見れば人間には一目で分かることが多くあります。

逆に、商品の外装やイメージだけのものなど、画像だけでは判断がつかない商品も存在します。テキストと画像は互いに補い合う関係にあり、両方をバランス良く扱えるようにしたい、というのがマルチモーダルモデルを使用する主な動機となります。

あわせて、分類先のカテゴリも見直しています。従来のモデルは100クラス程度の分類モデルであったのに対し、今回は階層構造を持つカテゴリマスタのフルパス（例: グルメ・飲料/スイーツ・お菓子/ケーキ）を1つのクラスとみなした 498クラスのシングルラベル分類にしています。

使用したモデル

マルチモーダルなモデルには様々な選択肢がありますが、今回は MMBT (MultiModal BiTransformer) を採用しました。

github.com

アイデアはとてもシンプルで、画像を数個の「単語」のようなトークンに変換して、テキストトークンと一緒に同じ1本のTransformerに流し込むというものです。画像とテキストをそれぞれ別のエンコーダに通して最後にベクトルを結合するのではなく、最初から同じself-attentionの中で混ぜてしまうアプローチになります。

_{https://arxiv.org/abs/1909.02950 より引用}

その際、attention内で個々のトークンがどのモーダルに対応するのかを、Token Type Embedding（セグメント埋め込み）によって区別しています。こちらは従来、テキスト文が質問と回答のどちらに該当するかなどのマーカーとして使用していたものですが、MMBTではそのトークンが属するモーダルによって値を変えたEmbeddingを位置埋め込みなどと同様に、それぞれのトークンに加算しています。

また、本実装では、論文のオリジナル実装から次の2点を置き換えています。

画像エンコーダ: ResNet → Swin Transformer (swin_base_patch4_window7_224.ms_in22k)
テキストエンコーダ: 英語BERT → 日本語BERT (tohoku-nlp/bert-base-japanese-v3)

実装は責務ごとに4つのクラスを積み上げる形になっています。以下、内側から順に見ていきます。

1. 画像を3つのトークンに変える

ImageEncoder は、224×224 の画像を Swin Transformer に通して 7×7 の特徴マップにし、それを3つの領域に平均プーリングして3本の特徴ベクトル（各1024次元）に要約します。この3本が「画像トークンの素」になります。

class ImageEncoder(nn.Module):
    def __init__(self, num_image_embeddings: int = 3):
        super().__init__()
        model = timm.create_model(
            "swin_base_patch4_window7_224.ms_in22k",
            pretrained=True,
            num_classes=0,
        )
        # avg poolingと最後のlinear層を除外
        modules = list(model.children())[:-2]
        self.model = nn.Sequential(*modules)
        self.pool = nn.AdaptiveAvgPool2d((num_image_embeddings, 1))

    def forward(self, x):
        out = self.model(x)
        out = out.permute(0, 3, 1, 2).contiguous()
        out = self.pool(out)
        out = torch.flatten(out, start_dim=2)
        return out.transpose(1, 2).contiguous()

画像を何トークンに要約するかはハイパーパラメータで、ここでは3としています。トークン数を増やすほど画像の情報は細かく残せますが、そのぶんテキストに使えるトークン列が短くなります。

2. 画像特徴を「BERTのトークン」に翻訳する

BERTのテキストトークンの埋め込みは、実際には 単語埋め込み + 位置埋め込み + セグメント埋め込み の足し算でできています。画像トークンもまったく同じ作り方にすることで、BERTは画像をテキストトークンと同様に扱うことができます。

ImageBertEmbeddings では、

Swin の出力（1024次元）を nn.Linear(1024, 768) でBERTの隠れ次元に射影し
位置埋め込みはテキスト側と同じものを共有し
画像かテキストかを区別するセグメント埋め込み（nn.Embedding(2, 768)）を新たに定義して加算する

という処理を行います。BERTの実装によってはそのまま使える token_type_embeddings が無いケースがあるため、こちらで定義しています。

さらに先頭に [CLS]、末尾に [SEP] の単語埋め込みを足すので、画像トークンの個数が3の場合、常に 固定長5トークン の並びになります。

[CLS] 画像1 画像2 画像3 [SEP] テキスト1 テキスト2 ... テキストN [PAD] ...
|___________________________| |_________________________________________|
   画像トークン（固定5個）            テキストトークン（可変長）

テキスト側は日本語BERTのtokenizerでトークン化し、最大長は 512 - 画像トークン数 に切り詰めます。

3. 1本のTransformerに流し込む

MultimodalBertEncoder は、画像トークン列とテキストトークン列を横に連結して1本のシーケンスにし、BERTのエンコーダに通します。

このとき attention mask は、画像部分は常に1（必ず全部見る）、テキスト部分は実トークンだけ1・パディングは0、という形で作っています。

attention_mask = torch.cat(
    [torch.ones(bsz, self.num_image_embeds + 2).long(), attention_mask],
    dim=1,
)
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)
# 1の部分は0に、0の部分は-10000.0に変換する
extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

あとは通常のBERTと同じで、全トークンが互いに attention を張り合い、最後に pooler が全体を768次元のベクトルに要約します。

4. 分類と埋め込みの取り出し

最後の MultimodalBertClf は、pooler の出力（768次元）を nn.Linear(768, クラス数) に通してカテゴリのスコアを出すだけの薄いクラスです。

class MultimodalBertClf(nn.Module):
    def __init__(self, n_classes, model, vocab, hidden_size: int = 768):
        super().__init__()
        self.enc = MultimodalBertEncoder(model=model, vocab=vocab)
        self.clf = nn.Linear(hidden_size, n_classes)

    def forward(self, txt, mask, segment, img):
        return self.clf(self.enc(txt, mask, segment, img))

    def embeddings(self, txt, mask, segment, img):
        # 分類前の特徴量を返す
        return self.enc(txt, mask, segment, img)

この768次元ベクトルは「商品テキストと商品画像の両方を要約したベクトル」なので、カテゴリ分類以外の下流タスクにも転用できます。推論時にはカテゴリと一緒にこの埋め込みも保存しています。

結果

学習はGeForce RTX 5090を積んだオンプレサーバーで数日程度行い、検証データ全体でのaccuracyは90%でした。推論バッチの作成後、新規登録された商品を対象に定性的な検証も実施したところ、498クラスのうち比較的少数のクラスに関しても、ある程度正確に推論できていました。また、従来のテキストのみによる分類では判別が難しかった商品についても、適切に分類できていることが確認できました。