시각 트랜스포머의 정보 흐름을 조절하는 변분 정보 병목

초록

본 논문은 비전 트랜스포머(ViT)의 어텐션이 잔차 스트림에 쓰는 정보를 변분 정보 병목(VIB)으로 제한함으로써, 통신 용량을 연속적으로 조절 가능한 하이퍼파라미터 β 하나만으로 독립 패치 처리부터 완전 글로벌 어텐션까지의 스펙트럼을 만든다. ImageNet‑100 실험을 통해 정확도와 라우팅된 정보량 사이의 트레이드오프를 정량화하고, 정보가 처음으로 흐르기 시작하는 어텐션 헤드와 패치‑레벨 업데이트 패턴을 분석한다. 제한된 내부 통신을 강제함으로써 모델의 메커니즘 해석이 용이해지고, 제어 가능성이 향상됨을 보인다.

상세 분석

이 연구는 비전 트랜스포머의 핵심 구조인 어텐션‑MLP 블록을 잔차 스트림(residual stream) 관점에서 재해석한다. 어텐션 블록은 패치 간 정보를 교환하는 유일한 경로이며, MLP 블록은 각 패치를 독립적으로 정제한다는 점에 착안해, 저자는 어텐션이 잔차에 쓰는 업데이트 Δℓi 직후에 변분 정보 병목을 삽입한다. 구체적으로, Δℓi 를 인코더 qϕ(·)를 통해 가우시안 잠재 변수 zℓi 로 매핑하고, 디코더 gθ(·)로 복원된 ˆΔℓi 를 잔차에 더한다. 이 과정에서 KL 발산 DKL(qϕ‖p) 를 손실에 β배 가중치로 추가함으로써, 각 헤드가 전달할 수 있는 평균 정보량을 직접 제어한다. β=0이면 모든 어텐션 업데이트가 차단돼 패치별 독립 투표(DeepSets와 유사)만 남고, β→∞이면 기존 ViT와 동일하게 자유로운 정보 흐름을 허용한다.

실험에서는 ViT‑tiny(12 블록, 3 헤드/블록)를 사용해 36개의 독립 병목을 동시에 학습시켰다. β 값을 로그 스케일로 변화시켜 만든 모델군은 정확도와 총 KL(즉, 전체 전송 정보량) 사이에 거의 선형에 가까운 파레토 프론트를 형성한다. 특히 중간 β 구간에서 정확도가 급격히 상승하면서도 전송된 정보량은 상대적으로 완만히 증가한다는 점이 흥미롭다. 이는 소량의 핵심 정보만으로도 글로벌 표현을 형성할 수 있음을 시사한다.

정보 라우팅 분석에서는 패치‑레벨 KL을 시각화해, 고β 모델에서도 전체 패치 중 소수(≈5~10%)만이 높은 정보량을 갖고 나머지는 거의 제로에 가깝다는 ‘희소 라우팅’ 현상을 발견했다. 또한, 어텐션 헤드별 활성화 패턴을 NMI(정규화된 상호 정보)로 정량화했을 때, 초기 레이어의 몇몇 헤드가 다른 헤드와 높은 상관성을 보이며, 이후 레이어로 갈수록 정보가 분산되고 독립적인 특성을 띤다. 이는 제한된 통신 용량 하에서 모델이 먼저 핵심적인 전역 관계를 포착하고, 이후에 세부적인 디테일을 점진적으로 추가한다는 메커니즘을 뒷받침한다.

또한, 저자는 KL 기반의 ‘패치 투표 다양성’ 지표(역심슨 지수)를 도입해 β가 증가함에 따라 패치별 로짓 범위와 클래스 다양성이 어떻게 변하는지 정량화했다. β가 낮을 때는 모든 패치가 동일한 클래스를 강하게 지지하지만, β가 높아지면 패치 간 로짓 차이가 커지고, 다양한 클래스가 동시에 높은 확률을 얻는다. 이는 어텐션이 정보를 공유함에 따라 개별 패치가 보다 풍부한 컨텍스트를 얻게 됨을 의미한다.

전반적으로 변분 정보 병목을 어텐션 업데이트에 직접 적용함으로써, 모델이 학습 단계부터 정보 흐름을 제한받아 보다 해석 가능하고 제어 가능한 내부 구조를 갖게 된다. 이는 기존의 사후 분석 방식과 달리, 최적화 과정 자체가 ‘정보 효율성’이라는 명시적 목표를 갖게 함으로써 메커니즘 해석 연구에 새로운 도구를 제공한다.