다이나믹 차등 선형 어텐션으로 고품질 이미지 생성 혁신

초록

본 논문은 선형 확산 트랜스포머(LiT)의 과도한 스무딩 문제를 해결하기 위해 동적 투영 모듈, 동적 측정 커널, 토큰 차등 연산자를 결합한 새로운 선형 어텐션 메커니즘인 DyDiLA를 제안한다. 이를 기반으로 만든 DyDi‑LiT는 기존 DiT와 최신 효율형 모델들을 능가하는 FID·sFID·KID·IS 등 다중 지표에서 우수한 성능을 보이며, 계산 복잡도는 O(N) 수준으로 유지한다.

상세 분석

DyDiLA는 선형 어텐션이 갖는 “저분산·과도한 스무딩” 현상을 세 가지 원인—토큰 이질성, 부적절한 유사도 측정, 컨텍스트 민감한 검색—에 대해 구조적으로 대응한다. 첫 번째로 제안된 동적 투영 모듈은 토큰별로 서로 다른 프로젝터를 라우팅함으로써, 시점·공간에 따라 달라지는 토큰 분포를 동일한 선형 변환에 강제로 매핑하는 것을 방지한다. 라우터는 학습 가능한 매트릭스 R_QP, R_KP와 토큰‑프로젝터 집합 {W_Qi, W_Ki}를 이용해 arg‑max 방식으로 최적 프로젝터를 선택한다. 이는 기존 Softmax 기반 DiT가 공유 프로젝터만을 사용해 토큰 표현을 과도하게 동질화시키는 문제를 완화한다.

두 번째인 동적 측정 커널은 기존 선형 어텐션이 모든 토큰에 동일한 커널 함수 ϕ(·)를 적용하는 한계를 넘어, 토큰별로 서로 다른 정규화 파라미터 γ_f를 갖는 커널 집합 {ϕ_f}를 라우팅한다. 토큰 Z_i는 라우터 R_ZF에 의해 가장 높은 스코어를 가진 커널로 매핑되어, 토큰의 방향성은 유지하면서도 의미적으로 유사한 토큰 간의 내적을 강화한다. 이는 ReLU‑파워 연산을 기반으로 한 norm‑preserving 변환을 확장한 것으로, 토큰 이질성을 고려한 정교한 유사도 측정을 가능하게 한다.

세 번째인 토큰 차등 연산자(TDO)는 Q와 K에 대한 “정보 중복” Q′, K′을 차감함으로써, 쿼리‑키 매칭의 강건성을 높인다. 차감 비율 λ_Q, λ_K는 각각 토큰‑쌍 (e_Q, e_Q′), (f_K, f_K′)에 대해 별도의 라우터(R_QD, R_KD)와 학습 가능한 스칼라 집합 {λ_d}를 통해 동적으로 선택된다. 최종 어텐션 출력은 (e_Q − λ_Q·e_Q′)·(f_K − λ_K·f_K′)^T·V 형태로 계산되며, 이는 기존 선형 어텐션이 갖는 “키 중복” 문제를 차등적으로 보정한다. 또한, 깊이별 3×3 depth‑wise convolution을 추가해 토큰 수준 차등 연산과 공간적 컨텍스트를 동시에 강화한다.

실험에서는 ImageNet‑1K와 그 서브셋(Sub‑IN)에서 256×256·512×512 해상도로 Small·Base·Large 모델을 학습하였다. DyDi‑LiT는 동일한 FLOPs(≈110 GFLOPs)에서 DiT‑S(Softmax)와 비교해 FID 11.81→7.15, sFID 0.0053→0.0080, IS 43.28→47.05 등 전반적인 지표가 크게 개선되었으며, 기존 선형 어텐션 기반 모델(Sana, DiG‑S 등)보다도 일관된 우위를 보였다. 특히, 동적 차등 연산을 attention‑map 수준으로 대체한 실험(ablation)에서는 토큰‑단위 차등이 더 높은 성능 향상을 제공함을 확인하였다. 계산 복잡도는 여전히 O(N)이며, 라우터와 추가 파라미터는 전체 파라미터 수의 2~3% 수준에 불과해 실용적이다.

한계점으로는 라우터 선택 과정이 arg‑max 기반이므로 미분 가능성에 제약이 있을 수 있으며, 토큰 수가 매우 큰 경우 라우터 매트릭스 자체가 메모리 부담을 초래할 가능성이 있다. 또한, 현재 실험은 주로 이미지 생성에 국한되어 있어 비디오·텍스트 등 다른 멀티모달 확산 모델에 대한 일반화 검증이 부족하다. 향후 연구에서는 소프트 라우팅(soft‑assignment)이나 메모리 효율적인 라우터 설계, 그리고 다양한 도메인에 대한 확장성을 탐색할 여지가 있다.