픽셀-공간 디퓨전 트랜스포머에서 레지스터의 중요성

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

디퓨전 모델이 트랜스포머 기반으로 전환되면서 픽셀 공간에서 고차원 목표를 직접 예측하는 모델들이 늘어나고 있다. 픽셀 공간은 잠재 공간보다 표현이 덜 구조화되어 중간 표현이 고주파 노이즈와 높은 토큰 노름을 보이기 쉬운데 이는 생성물의 구조적 일관성을 저해했다. 본 논문은 추가 학습 토큰인 레지스터가 이러한 문제를 완화하며 특히 픽셀-공간 DiT에서 생성 품질을 실질적으로 개선한다는 점을 규명해 모델 설계와 조건부 토큰 사용 방식에 실무적 의미를 제공한다.

왜 중요한가

핵심 기여

Diffusion Transformer는 ViT와 달리 패치 토큰 고노름 아웃라이어를 보이지 않음

여러 크기와 공간 설정에서 DiT는 DINOv2 같은 ViT와 대비해 패치 토큰의 노름 분포가 균일하게 나타났다. 이 관찰은 DiT가 원래 레지스터 도입의 주요 명분인 패치 아웃라이어를 갖지 않음을 의미한다. 그럼에도 불구하고 레지스터를 추가하면 고노름이 레지스터 자체로 옮겨지는 현상이 확인되었다.

레지스터는 픽셀-공간 DiT에서 중간 표현을 정리하고 생성 품질을 향상시킴

레지스터를 사용하면 모든 패치 토큰의 특징 노름이 감소하고 Total Variation이 낮아져 중간 표현이 더 공간적으로 매끄러워졌다. 이 효과는 노이즈가 큰 초기 확산 단계에서 특히 뚜렷했으며 그 결과 FID가 일관되게 개선되는 경향이 관찰되었다. 같은 실험을 잠재 공간 모델에 적용하면 개선폭이 작거나 손해가 발생해 효과가 공간 의존적임이 드러났다.

레지스터의 이중 역할 규명 및 Register Guidance 기법 제시

선형 프로빙과 주의맵 시각화를 통해 일부 레지스터는 전역 의미를 캡처하는 반면 다른 레지스터는 노름 싱크로 특화된다는 점이 확인되었다. 이러한 특성에 기반해 원래 레지스터가 없는 모델을 약한 모델로 삼아 레지스터가 캡처한 방향을 증폭하는 Register Guidance 기법을 도입했다. Register Guidance를 CFG와 조합하면 최근 픽셀-공간 DiT에서 일관된 성능 향상이 관찰되었다.

핵심 아이디어 이해하기

이미지 생성용 트랜스포머는 입력 이미지 패치를 시퀀스로 취급해 self-attention으로 정보를 결합한다. 이때 각 패치가 고차원 픽셀 목표를 예측하도록 학습되면 모든 패치 표현에 많은 정보량과 변동성이 실리며 결과적으로 특징 벡터의 노름이 전체적으로 커지고 공간적 잡음이 증가할 수 있다. 이러한 상태에서는 일부 토큰이 과도한 노름을 갖는 ViT 유형의 아웃라이어가 나타나기도 하지만, DiT에서는 손실이 모든 패치에 균등하게 부과되어 패치가 높은 노름을 지속하기 어렵다.

방법론

연구는 픽셀-공간 pDiT와 PixelDiT, 잠재 공간의 SiT 및 RAE 등 여러 백본과 해상도 256×256, 패치 크기 16 설정에서 실험을 수행했다. 레지스터는 학습 가능한 추가 토큰으로 패치 시퀀스에 덧붙여지며 이 토큰들은 학습 손실 계산에 포함되지 않는다. 분석 기법으로는 토큰별 특징 노름 측정, Total Variation 기반의 중간 표현 공간적 매끄러움 측정, 선형 프로빙을 통한 전역 의미 평가, 그리고 레지스터 수와 적용 시작 레이어에 대한 어블레이션이 활용되었다.

주요 결과

표와 그림들을 통해 레지스터를 도입한 픽셀-공간 모델에서 FID가 일관되게 개선되는 결과가 보고되었다. 선형 프로빙 결과에서는 일부 레지스터 토큰이 약 0.9 수준의 높은 분류 정확도를 보인 반면 노름이 가장 큰 레지스터는 약 0.02 수준의 낮은 정확도를 기록해 노름 싱크와 의미 인코딩의 이중 역할이 수치로 확인되었다. Total Variation 비교에서는 레지스터 적용 모델이 낮은 확산 시점 영역에서 TV 비율이 1보다 작아 중간 표현이 더 매끄럽다는 정량적 근거를 제공했다.

기술 상세

아키텍처 관점에서 레지스터는 패치 시퀀스에 학습 가능한 토큰으로 추가되며 이 토큰들은 모델의 출력 손실에 포함되지 않는다. 실험에서는 pDiT-B/16 등 여러 크기 모델과 PixelDiT 구성을 사용했고 학습은 flow matching으로 ImageNet 256×256에서 수행되었다. 레지스터의 영향은 주로 MLP 층 이후에 형성되는 경향이 확인되어 내부 연산 흐름에서 활성화 크기가 변하는 지점을 파악했다.

한계점

저자들은 레지스터 효과가 공간 의존적이며 픽셀-공간에서 가장 강하게 관찰된다고 명시했다. 잠재 공간에서는 개선폭이 작거나 무시할 수준이며 RAE 기반 모델에서는 성능 저하가 보고되어 모든 백본에 보편적으로 적용 가능한 해결책은 아니었다. 또한 어떤 경우에는 텍스트-투-이미지 대형 모델에서 패치 아웃라이어가 불안정하게 나타나 어휘적 일반화에 제약이 있음을 지적했다.

실무 활용

레지스터 토큰은 픽셀-공간 DiT의 표현을 정규화해 구조적 일관성과 생성 품질을 향상시키는 실무적 수단으로 활용 가능하다. 논문에서 사용한 코드 저장소가 공개되어 있으므로 동일한 백본에 쉽게 레지스터를 추가해 실험할 수 있다. 다만 효능은 학습 공간과 레지스터 위치에 따라 달라지므로 도입 전 레이어 시작점과 토큰 수에 대한 어블레이션을 권장한다.

고해상도 픽셀-공간 이미지 생성 파이프라인에서 레지스터를 추가해 중간 표현의 공간적 노이즈를 줄이고 FID 개선을 노릴 수 있다. 실험은 ImageNet 256×256 기준으로 수행되었으며 flow matching 학습으로 성능을 검증했다. 레지스터는 깊은 블록에서 시작하는 설정이 보다 안정적인 성능 개선을 보였다.
조건부 생성에서 in-context 클래스 토큰을 레지스터와 유사한 방식으로 배치해 초기 표현을 더 잘 형성하도록 설계할 수 있다. 본문은 in-context 토큰이 암묵적 레지스터 역할을 수행하며 일부 성능 향상은 토큰의 레지스터 성질에 기인함을 확인했다. 따라서 클래스나 텍스트 조건을 도입할 때 토큰 수와 삽입 레이어를 조정하면 추가 이득을 얻을 수 있다.
이미 존재하는 DiT 기반 시스템에 Register Guidance를 적용해 레지스터가 포착한 방향을 증폭하는 방식으로 후처리 제어를 강화할 수 있다. 논문은 Register Guidance를 CFG와 결합해 추가적인 품질 향상을 확인했으며 공개된 코드로 재현이 가능하다. 실제 적용 시에는 레지스터 유무의 비교 실험과 샘플링 단계에서의 파라미터 튜닝이 필요하다.

코드 공개 여부: 공개

코드 저장소 보기

키워드

용어 해설

Register Token: — 추가로 학습되는 특수 토큰으로 패치 토큰과 함께 트랜스포머 입력에 결합된다. 이 토큰은 손실에 직접 참여하지 않아 높은 노름 값을 흡수하거나 전역 정보를 집약하는 용도로 사용된다. 본문에서는 DiT에서 특이값 처리를 위한 슬롯으로 동작하며 패치 표현의 노이즈를 줄이는 역할을 수행한다.
Total Variation: — 중간 표현의 공간적 부드러움을 정량화하는 지표로 인접 픽셀 간 값 차이를 합산해 계산된다. 값이 작을수록 특징 맵이 더 매끄럽고 큰 값은 고주파 노이즈가 많음을 의미한다. 논문에서는 register가 적용된 모델에서 TV가 낮아지는지를 비교해 고잡음 단계에서 표현이 정리되는지를 평가했다.
Flow Matching: — 확률적 생성과정의 연속성을 직접 모델링하는 학습 방식으로 확률 흐름을 맞추는 손실을 사용한다. 이미지 생성에서 노이즈 레벨별 목표 벡터장을 예측하도록 학습하여 샘플링을 수행한다. 본 연구의 pDiT 모델들은 ImageNet 256×256에서 flow matching으로 학습되었다.
Pixel Space: — 이미지의 원시 픽셀 값으로 직접 학습하는 공간을 의미하며 오토인코더 기반 잠재 공간과 대비된다. 고차원이고 구조화가 덜 되어 있어 중간 표현에 높은 노이즈와 큰 노름이 발생하기 쉽다. 본문에서는 픽셀 공간에서 register의 이득이 특히 뚜렷하게 나타남이 보고되었다.
Linear Probing: — 학습된 표현의 정보량을 평가하기 위해 고정된 특징 위에 간단한 선형 분류기를 학습하는 방식이다. 특징이 얼마나 전역적 의미나 클래스 정보를 담고 있는지 측정하는 데 사용된다. 본문에서는 register 토큰 특징을 선형 프로빙해 일부 토큰이 높은 분류 정확도와 낮은 노름을 갖는 반면 다른 토큰은 노름만 큰 규격을 보였음을 확인했다.