Concept Bottleneck Models

개념 병목 모델(CBM)은 입력에서 사람이 이름을 붙인 개념 집합을 먼저 예측한 다음, 해당 개념만 사용해 최종 레이블을 예측하는 신경망입니다. 무릎 엑스레이에서는 관절 간격 협소화와 골극이, 새 사진에서는 날개 색과 부리 모양이 개념이 될 수 있습니다. 모든 예측이 레이블이 지정된 이 좁은 층을 통과하므로, 사람은 어떤 개념이 판단을 이끌었는지 확인하고 잘못된 개념을 수정해 결과를 바꿀 수 있습니다. 따라서 CBM은 아키텍처에 내장된 설명 가능한 AI(XAI)의 한 형태입니다. Stanford와 Google의 Pang Wei Koh, Thao Nguyen, Yew Siang Tang, Stephen Mussmann, Emma Pierson, Been Kim, Percy Liang은 ICML 2020에서 개념 병목 모델이라는 용어를 도입했습니다. 이들은 엑스레이에서 무릎 골관절염의 등급을 매기는 작업과 조류 종을 식별하는 작업에서 표준 종단 간 모델에 필적하는 정확도를 보였습니다.

개념 병목 모델의 작동 방식#

CBM은 두 부분으로 구성됩니다. 입력(예: 이미지)을 개념 값에 매핑하는 개념 예측기와, 개념 값을 출력에 매핑하는 레이블 예측기입니다. Koh와 동료들은 두 부분을 학습하는 세 가지 방법을 비교했습니다.

• 독립 방식: 개념 예측기는 개념 주석을 사용하고 레이블 예측기는 실제 개념을 사용해 각각 별도로 학습합니다.
• 순차 방식: 개념 예측기를 먼저 학습한 다음, 레이블 예측기는 예측된 개념을 사용해 학습합니다.
• 공동 방식: 두 부분을 개념 손실과 레이블 손실의 가중합으로 함께 학습합니다.

켈그렌-로렌스 중증도 등급이 레이블로 지정된 무릎 엑스레이 36,369건이 포함된 Osteoarthritis Initiative 데이터에서는 방사선 전문의가 이미 평가하는 관절 간격 협소화, 골극, 골경화 등의 순서형 소견 10개를 개념으로 사용했습니다. 200종의 새 사진 11,788장이 포함된 Caltech-UCSD Birds 데이터셋에서는 날개 색과 부리 모양 등의 이진 속성 112개를 사용했습니다. 모든 학습 예제에는 개념과 레이블 주석이 모두 필요하며, 이것이 이 접근 방식의 주요 비용입니다.

병목 구조는 테스트 시 개입을 가능하게 합니다. 전문가가 예측된 개념이 잘못되었음을 확인하면 올바른 값으로 바꾸고, 레이블 예측기가 출력을 다시 계산합니다. 엑스레이 작업에서는 개념 두 개의 실제 값만 입력해도 예측된 중증도 등급의 오차가 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개의 레이블을 올바르게 예측합니다. 예제마다 확신도가 가장 낮은 개념 한 개, 두 개, 세 개의 실제 값을 입력하면 두 모델을 변경하지 않고도 레이블 정확도가 각각 0.861, 0.915, 0.951로 높아집니다.

CBM의 활용 분야#

CBM은 전문가가 이미 이름이 지정된 소견을 바탕으로 추론하는 분야에 적합합니다. 의료 영상 분석에서는 앞서 언급한 엑스레이 연구가 방사선 전문의가 기록하는 것과 같은 소견을 사용하므로, 임상의가 각 소견을 확인할 수 있습니다. 안전 모니터링도 같은 구조입니다. Construction-PPE 데이터셋으로 학습한 Ultralytics YOLO26 모델은 helmet, no_helmet, vest, no_boots과 같이 이름이 지정된 클래스를 감지합니다. 이러한 감지 결과만 읽는 규정 준수 규칙은 개념 병목이 되어, 검토자가 특정 보호 장비의 누락 여부를 근거로 판단을 추적할 수 있습니다.

후속 연구에서는 주석 요건을 완화했습니다. 레이블 없는 CBM은 언어 모델로 개념 목록을 생성하고 CLIP으로 개념을 평가해 개념 레이블 없이도 이 접근 방식을 ImageNet 규모로 확장합니다. 사후 CBM은 임베딩을 개념 방향에 투영해 이미 학습된 모델을 변환합니다.

한계#

개념 집합은 레이블을 결정하는 모든 요소를 포착하는 경우가 드물어 CBM의 정확도에 상한을 둡니다. 레이블 예측기가 명확한 값 대신 개념 확률을 받으면, 각 개념의 의미를 넘어서는 레이블 관련 정보가 해당 확률에 담길 수 있습니다. 개념 누출이라고 하는 이 현상은 설명과 개입의 신뢰성을 떨어뜨립니다. 개념 임베딩 모델은 개입을 지원하면서 정확도를 높게 유지하기 위해 각 개념을 학습된 벡터로 표현합니다. 정의가 불명확하거나 서로 상관된 개념은 신뢰할 수 없는 설명을 만들며, 편향된 개념 주석은 이에 의존하는 모든 판단에 편향을 전달합니다.

CBM과 관련 개념 비교#

• 돌출 맵: 이미 학습된 블랙박스 모델에 영향을 준 픽셀을 강조합니다. CBM은 이름이 지정된 개념으로 판단을 설명하고 해당 개념을 수정할 수 있게 합니다.
• 기계론적 해석 가능성: 학습된 모델이 스스로 익힌 특징과 회로를 역설계합니다. CBM은 지도 학습을 통해 중간 개념을 미리 고정합니다.
• 선형 프로빙: 모델의 예측 방식을 바꾸지 않고 학습된 모델의 활성화에서 개념을 디코딩할 수 있는지 확인합니다. Been Kim과 동료들의 TCAV와 같은 블랙박스 개념 프로브는 사후에 개념의 영향을 측정합니다.
• 특징 공학: 원시 데이터에서 입력 특징을 수작업으로 계산합니다. CBM의 개념은 학습된 모델이 원시 입력에서 예측하고 사람의 레이블로 지도 학습합니다.