Concept Bottleneck Models

コンセプトボトルネックモデル（CBM）は、入力から人が名前を付けた概念をまず予測し、その概念だけを使って最終ラベルを予測するニューラルネットワークです。膝のX線画像では関節裂隙の狭小化や骨棘、鳥の写真では翼の色やくちばしの形が概念になります。すべての予測がこの狭いラベル付きレイヤーを通るため、どの概念が判断に影響したかを人が確認でき、誤った概念を修正して結果を変えられます。そのため、CBMはアーキテクチャに組み込まれた説明可能なAI（XAI）の一種です。StanfordとGoogleのPang Wei Koh、Thao Nguyen、Yew Siang Tang、Stephen Mussmann、Emma Pierson、Been Kim、Percy Liangは、ICML 2020でコンセプトボトルネックモデルを発表してこの用語を導入し、膝のX線画像から変形性膝関節症の重症度を判定するタスクと鳥の種を識別するタスクの2つで、標準的なエンドツーエンドモデルに匹敵する精度を示しました。

概念ボトルネックモデルの仕組み#

CBMは2つの部分で構成されます。画像などの入力を概念値に写像する概念予測器と、概念値を出力に写像するラベル予測器です。Kohらは、両者を学習する3つの方法を比較しました。

• 独立型： 概念予測器は概念アノテーションから、ラベル予測器は真の概念から、それぞれ別々に学習します。
• 逐次型： まず概念予測器を学習し、その予測概念を使ってラベル予測器を学習します。
• 同時型： 概念損失とラベル損失の加重和を用いて、両方の部分を同時に学習します。

重症度をKellgren-Lawrenceグレードでラベル付けした36,369枚の膝X線画像を含むOsteoarthritis Initiativeデータでは、概念として、関節裂隙の狭小化、骨棘、骨硬化など、放射線科医がすでに評価している10個の順序尺度所見を用いました。200種の鳥の写真11,788枚を含むCaltech-UCSD Birdsデータセットでは、翼の色やくちばしの形など112個の二値属性を用いました。すべての学習例に概念とラベルの両方のアノテーションが必要であり、これがこの手法の主なコストです。

ボトルネックにより、テスト時の介入が可能になります。専門家が予測された概念に誤りがあると判断した場合、正しい値に置き換えると、ラベル予測器が出力を再計算します。X線画像のタスクでは、わずか2つの概念に真の値を与えるだけで、予測された重症度グレードの誤差が0.4超からRMSE約0.3に低下しました。

このPyTorchの例では、6つの二値概念がノイズを含む入力を生成し、ラベルが概念だけに依存するおもちゃのデータで独立型CBMを学習します。その後、専門家が確信度の最も低い概念に真の値を与えるときのラベル精度を測定します。

import torch
from torch import nn

torch.manual_seed(0)
# Toy data: 6 binary concepts generate noisy 20-D inputs, and the label depends only on the concepts
concepts = torch.randint(0, 2, (3000, 6)).float()
inputs = concepts @ torch.randn(6, 20) + 2 * torch.randn(3000, 20)
labels = (concepts @ torch.tensor([3.0, 2, 2, 1, 1, 1]) >= 5).float()
train, test = slice(0, 2000), slice(2000, None)

# Independent training: inputs -> concepts, and concepts -> label, each against its own ground truth
concept_model, label_model = nn.Linear(20, 6), nn.Linear(6, 1)
optimizer = torch.optim.Adam([*concept_model.parameters(), *label_model.parameters()], lr=0.05)
bce = nn.functional.binary_cross_entropy_with_logits
for _ in range(500):
    optimizer.zero_grad()
    loss = bce(concept_model(inputs[train]), concepts[train]) + bce(label_model(concepts[train])[:, 0], labels[train])
    loss.backward()
    optimizer.step()

with torch.no_grad():
    probs = concept_model(inputs[test]).sigmoid()
    print(f"Concept accuracy: {((probs > 0.5) == concepts[test]).float().mean():.3f}")
    for k in range(4):
        # Test-time intervention: replace the k least certain predicted concepts with their true values
        uncertain = (probs - 0.5).abs().argsort(dim=1)[:, :k]
        fixed = probs.scatter(1, uncertain, concepts[test].gather(1, uncertain))
        accuracy = ((label_model(fixed)[:, 0] > 0) == labels[test].bool()).float().mean()
        print(f"Intervene on {k} of 6 concepts: label accuracy {accuracy:.3f}")

概念予測器はテスト用概念6,000件のうち0.834件を正しく予測し、CBMはテスト例1,000件のうち0.819件に正しいラベルを付けます。各例で確信度が最も低い概念を1つ、2つ、3つ選び、それぞれに真の値を与えると、ラベル精度は0.861、0.915、0.951に上昇します。どちらのモデルも変更しません。

CBMの活用分野#

CBMは、専門家がすでに名前付きの所見を使って判断する分野に適しています。医用画像解析では、前述のX線画像研究で放射線科医が記録するものと同じ所見を用いたため、臨床医が各所見を確認できます。安全監視も同様の構造です。Construction-PPEデータセットで学習したUltralytics YOLO26モデルは、helmet、no_helmet、vest、no_bootsなどの名前付きクラスを検出します。これらの検出結果だけを参照するコンプライアンスルールは、CBMとして機能し、レビュー担当者は不足している項目をたどって判断を確認できます。

後続の研究では、アノテーション要件が緩和されました。ラベル不要のCBMは言語モデルで概念リストを生成し、CLIPで概念をスコアリングすることで、概念ラベルなしでImageNetにまで適用範囲を広げます。また、事後CBMは、学習済みモデルの埋め込みを概念方向に射影して、既存モデルを変換します。

制限事項#

概念セットだけではラベルを決める要因をすべて捉えられないことが多く、CBMの精度には上限があります。ラベル予測器が確定値ではなく概念の確率を受け取る場合、その確率に各概念の意味を超えたラベル情報が含まれることがあります。これは概念漏洩と呼ばれ、説明や介入の信頼性を損ないます。Concept Embedding Modelsは、精度を高く保ちながら介入にも対応できるよう、各概念を学習可能なベクトルで表現します。定義が不適切な概念や相関した概念は信頼性の低い説明につながり、偏った概念アノテーションは、それに依存するすべての判断にバイアスを持ち込みます。

CBMと関連概念の比較#

• 顕著性マップ： 学習済みのブラックボックスモデルに影響したピクセルを強調表示します。CBMは名前付きの概念で判断を説明し、その概念を編集できます。
• メカニスティック解釈可能性： 学習済みモデルが独自に学んだ特徴や回路をリバースエンジニアリングします。CBMでは教師あり学習によって、中間概念を事前に固定します。
• 線形プロービング： 予測方法を変えずに、学習済みモデルの活性化から概念を復号できるかを検証します。Been KimらによるTCAVのような、ブラックボックスモデルに対する概念プローブは、事後的に概念の影響を測定します。
• 特徴量エンジニアリング： 生データから入力特徴量を手作業で計算します。CBMの概念は、学習済みモデルが生の入力から予測し、人手によるラベルで教師あり学習します。