DINO 자기지도학습 원리 — 비전 트랜스포머를 라벨 없이 훈련하기

핵심 메시지

• DINO(distillation with no labels)는 라벨 없는 이미지 데이터만으로 비전 트랜스포머(ViT)를 학습시키는 자기지도학습 방법이다.
• BERT·GPT가 라벨 없는 텍스트로 사전학습된 것처럼, 이미지에도 학생-교사(student-teacher) 구조를 적용해 라벨 없이 학습한다.
• 하나의 이미지에서 국소·전역 크롭을 여러 개 만들고(멀티크롭), 학생의 분포가 교사의 분포를 따라가도록 교차 엔트로피로 학습한다.
• 교사는 역전파로 학습하지 않고 학생의 지수이동평균(m≈0.996)으로 천천히 갱신되어, 안정적인 학습 목표 역할을 한다.
• 분할을 따로 가르치지 않아도 네트워크가 객체를 분할하는 능력과, 시각 검색에 쓸 수 있는 풍부한 이미지 임베딩을 스스로 학습한다.

쉽게 이해하기

2020년 등장한 비전 트랜스포머(ViT)는 2017년 '어텐션'가 논문의 구조를 이미지에 적용해, 이미지를 패치로 나누고 위치 임베딩과 분류용 CLS 토큰을 더해 인코더에 통과시킨 뒤 CLS 벡터로 분류를 수행한다. 다만 ViT는 합성곱 신경망(CNN)보다 데이터 효율이 낮아, 수억 장 규모의 라벨 데이터가 있어야 CNN을 앞선다. 문제는 그런 라벨 데이터를 대규모로 모으기가 매우 어렵다는 점이다.

그래서 2021년 페이스북 연구진은 '라벨 없는 대규모 이미지로 ViT를 학습시킬 수 있을까'를 물었다. 당시 BERT(마스킹된 언어 모델링, 다음 문장 예측)와 GPT(언어 모델링)가 라벨 없는 텍스트에서 스스로 학습 신호를 만들어 사전학습되고 있었다. DINO는 이 자기지도학습을 이미지에 적용하되, 이미지에는 텍스트처럼 손쉬운 라벨이 없으므로 학생-교사 구조를 도입한다.

핵심은 멀티크롭이다. 한 이미지에서 전체의 5~50%를 잘라낸 국소 크롭과 50% 이상을 담은 전역 크롭을 만들어, 학생에게는 국소 8개와 전역 2개를, 교사에게는 전역 크롭만 준다. 학생과 교사는 처음에 완전히 동일한 구조·파라미터로 시작한다. 각 크롭의 CLS 토큰 출력을 MLP를 거쳐 약 6만 5천 차원(라벨이 없으므로 '유사 객체'를 담을 여유 공간)으로 확장하고 온도 소프트맥스로 확률 분포를 만든 뒤, 학생 분포가 교사 분포를 따라가도록 교차 엔트로피 손실로 학습한다.

안정적 학습을 위한 장치가 여럿 있다. 온도 소프트맥스에서 학생은 온도가 높아(덜 뾰족해) 기울기 전파가 잘 되고, 교사는 온도가 낮아(더 뾰족해) 분명한 정답 역할을 한다. 교사가 지나치게 뾰족해지면 소수 뉴런만 활성화되는 모드 붕괴가 생기는데, 활성값의 평균을 빼는 센터링으로 이를 막는다. 또 교사에는 기울기를 전파하지 않고(stop gradient), 학생 파라미터의 지수이동평균(m≈0.996)으로 천천히 갱신해 목표가 흔들리지 않게 한다.

추론 시에는 학생이든 교사든 상관없다. 오래 학습하면 둘의 가중치가 사실상 같아지기 때문이다. 흥미로운 발견은 두 가지다. 첫째, 분할을 명시적으로 가르치지 않았는데도 CLS 어텐션 벡터를 시각화하면 객체 분할 지도가 나타난다. 학생이 귀·코·눈 같은 국소 크롭을 모두 하나의 개 객체로 묶도록 강제되면서 분할 능력이 저절로 창발한 것이다. 둘째, CLS 벡터 자체가 이미지의 풍부한 표현이 되어 k-최근접 이웃 기반 시각 검색에 활용된다.

주요 인사이트

• 라벨 수집이 병목일 때, 데이터 자체에서 학습 신호를 만드는 자기지도학습은 텍스트뿐 아니라 이미지에도 강력하게 적용된다.
• 학생-교사 구조의 목적은 모델 축소나 추론 가속이 아니라 '안정적인 학습 목표'를 만드는 데 있다.
• 교사를 학생만큼 빠르게 갱신하면 목표가 계속 움직여 학습이 무너지므로, 느린 이동평균으로 안정성을 확보하는 것이 관건이다.
• 센터링과 온도 소프트맥스의 조합은 모드 붕괴를 막아 6만 5천 개의 출력 뉴런이 골고루 쓰이게 한다.
• 명시적 지도 없이도 멀티크롭이라는 학습 방식만으로 객체 분할과 검색용 임베딩이 창발한다는 점이 DINO의 핵심 통찰이다.

자주 묻는 질문

관련 AI 소식