애플, 확산 트랜스포머 속 '이상 토큰' 문제 원인 규명

무엇이 나왔나

애플 ML 리서치가 이미지 생성에 쓰이는 확산 트랜스포머(Diffusion Transformer, DiT) 내부에서 벌어지는 '이상 토큰(outlier token)' 현상을 정면으로 다룬 논문을 8월 공개했다. 연구팀은 대표성 오토인코더(RAE) 기반 DiT 파이프라인에서 사전학습된 비전 트랜스포머(ViT) 인코더와 DiT 자체 모두에서 이 현상이 나타난다는 점을 확인했다. 특히 DiT 내부의 중간 층에서 이런 토큰이 두드러졌다.

연구팀이 발견한 흥미로운 점은, 값이 큰 이상 토큰을 그냥 가리거나 지우는 방식으로는 성능이 좋아지지 않았다는 것이다. 문제가 몇 개의 극단값 때문이 아니라, 해당 위치의 지역 패치 정보 자체가 손상돼 있기 때문이라는 뜻이다. 이를 해결하려고 제안한 것이 '이중 단계 레지스터(Dual-Stage Registers, DSR)'다. 학습된 레지스터를 쓸 수 있으면 그것을 활용하고, 없으면 추론 시점에 반복적으로 적용하는 레지스터를 쓰고, 디노이저(잡음 제거 네트워크)에는 별도의 확산 레지스터를 붙이는 방식이다.

왜 '이상 토큰'이 문제인가

비전 트랜스포머는 이미지를 작은 조각(패치)으로 잘라 각각을 하나의 '토큰'처럼 다룬다. 문장을 단어 단위로 쪼개 처리하는 언어모델과 같은 방식이다. 그런데 학습이 진행되면서 일부 토큰이 이상할 정도로 큰 값(고노름, high-norm)을 갖고, 정작 그 위치의 실제 이미지 정보는 거의 담지 못한 채 다른 토큰들의 주의(attention)를 과도하게 끌어가는 현상이 보고돼 왔다. ViT 연구에서는 이런 토큰에 별도의 '레지스터' 자리를 마련해 흡수시키는 방법이 알려져 있었지만, 생성모델인 DiT에서 이 현상이 어떻게 작동하는지는 그동안 제대로 규명되지 않았다.

DiT는 스테이블 디퓨전 3, FLUX 계열 등 최신 이미지 생성모델이 기존의 U-Net 구조 대신 채택하고 있는 아키텍처다. 애플은 앞서 확산 기반 방식과 자기회귀 방식의 성능을 비교하거나, 확산 언어모델을 17억 파라미터로 키우는 실험 등 확산모델 계열 연구를 이어온 바 있다. 이번 논문은 그 연장선에서, 이미지 품질을 갉아먹는 미세한 결함의 근본 원인을 짚어낸 결과로 볼 수 있다.

그래서 뭐가 달라지나

이상 토큰 문제는 그동안 생성 이미지에서 이따금 나타나는 국소적 얼룩이나 왜곡된 패치의 원인 중 하나로 지목돼 왔다. 애플의 이번 연구는 이를 가리는 임시방편이 아니라 학습 구조 차원에서 다루는 접근을 제시했다는 점에서, 향후 DiT 기반 이미지·영상 생성모델의 안정성을 높이는 설계 지침으로 이어질 가능성이 있다. 다만 이번 결과는 연구 논문 수준의 검증이며, 실제 상용 모델에 어떤 형태로 반영될지는 추후 공개될 후속 연구를 지켜봐야 한다.