확산 모델 토큰 병합을 디노이징 동역학으로 결정하는 DaTMe

무엇인가

확산 모델은 반복 디노이징 과정에서 매 스텝마다 대량의 토큰을 다시 계산하기 때문에 추론 비용이 크다. 토큰 병합은 비슷한 토큰을 대표 토큰으로 합쳐 이후 레이어 연산을 줄이는 학습 불필요 가속 전략인데, 핵심 난제는 어떤 토큰을 흡수시켜도 안전하고 어떤 토큰은 앵커로 남겨야 하는지 정하는 일이다. ToMeSD 같은 초기 방식은 내용과 무관하게 지역 후보군을 만든 뒤 특징 유사도만으로 병합 쌍을 정해서, 내용상 중요한 토큰이 병합 소스로 뽑혀 세부 정보가 뭉개지는 문제가 있었다. 후속 연구는 classifier-free guidance(CFG)의 조건부·비조건부 노이즈 예측 차이를 중요도 신호로 써서 텍스트가 강하게 조종하는 토큰을 보존했지만, 저자들은 이것이 프롬프트 중심 편향을 갖는다고 지적한다.

어떻게 동작하나

논문은 CFG 신호의 실패를 두 갈래로 분석한다. 첫째, CFG는 시간에 따라 변하는 토큰 중요도를 포착하지 못한다. 확산 생성은 거친 구조에서 세부로 진행되는데, 스텝 간 변화가 큰 토큰은 모델이 실제로 디테일을 주입하는 영역이다. 이런 토큰의 공간 분포는 초반에는 전경에 몰려 있다가 후반 스텝에서 배경을 포함한 넓은 영역으로 퍼지는데, CFG 상위 토큰은 전 구간 전경에 고정돼 있어 정작 보호해야 할 후반부 배경 토큰을 놓친다. 둘째, CFG는 프롬프트와 무관한 구조적 중요도에 둔감하다. 프롬프트에 명시되지 않은 배경 물체, 균질 영역 사이의 경계, 미세 텍스처는 텍스트 유도에 거의 기여하지 않아 모든 타임스텝에서 낮은 점수를 받고 이웃한 균질 영역으로 병합되기 쉽다.

무엇과 다른가

제안 방법 DaTMe는 Tweedie denoised estimate를 기반으로 토큰 중요도를 두 축으로 정의한다. 시간 축은 예측된 clean image 공간에서의 스텝 간 잔차 |x̂0(t) − x̂0(t−1)|다. 저자들은 원시 score slope |s(t) − s(t−1)|를 차분 항과 스케일 항으로 분해해, 스케일 항이 사전 정의된 노이즈 스케줄에서 생기며 σt→0인 후반 스텝에서 급격히 커져 신호를 오염시킨다는 것을 보인다. Tweedie 차분은 이 스케일 항을 제거하고 차분 항만 남기며, σt−1/αt−1 계수로 후반에 자연 감쇠한다. 실증으로도 CFG 중요도는 전 구간 자기상관이 높아 정적이고, 원시 score slope는 정규화 타임스텝 0.4~0.6 부근에서 최종 생성 결과와의 상관이 급락하는 반면, Tweedie 차분은 더 높은 상관과 더 낮은 자기상관을 유지했다.

어떻게 쓰나

공간 축은 Tweedie 추정값과 그 coarse-scale 근사(다운샘플 후 업샘플, s=4)의 잔차로 정의한다. 이는 고주파 통과 필터처럼 작동해 윤곽, 경계, 미세 텍스처에 높은 응답을 주고 국소적으로 균질한 영역에는 낮은 응답을 준다. 이 두 축 위에서 DaTMe는 주어진 압축률 r에 대해 토큰 역할을 적응적으로 배정한다. 두 축 모두에서 하위 r 비율에 든 토큰의 교집합을 소스 집합으로, 두 축 모두에서 상위 (1−r) 비율에 든 토큰의 교집합을 제약 집합으로 삼아 구조적 앵커로 보존한다. 교집합 특성상 소스 집합이 목표 예산보다 작으므로, 토큰 시퀀스를 공간 특징 맵으로 바꿔 P=4개의 격자 셀에서 균등 샘플링해 소스를 보충한다. 남은 토큰은 목적지 풀이 되고, DiT가 공간 귀납 편향을 갖지 않는다는 점을 감안해 특징 유사도가 아니라 공간 인접성을 기준으로 이분 매칭을 수행한다.

전제와 한계

실험은 PixArt-α에서 COCO 2014 검증셋에서 뽑은 1만 개 이미지-프롬프트 쌍으로 진행했고, 압축률 0.3, 0.5, 0.7을 평가했다. DaTMe는 DiT 백본의 중간 범위 self-attention 레이어에 적용했으며, 확장성 검증을 위해 FLUX.1-dev의 joint-attention 블록 중간부터 최종 블록까지 같은 하이퍼파라미터로 적용했다. 지표는 CLIP-T, FID, HPSv2, SSIM, PSNR, LPIPS, Aesthetic Predictor V2.5 기반 미학 점수다. 정성적으로 ToMeSD는 공간 구조를 심하게 훼손했고, CFG 기반 방식은 프롬프트에 명시된 전경은 지켰지만 배경과 약하게 조건화된 영역 보존이 눈에 띄게 약했다. FLUX에서는 CFG를 쓰지 않아 CFG 기반 토큰 선택 자체가 적용 불가능한 반면 DaTMe는 그대로 적용됐다.

원문에 제시된 수치는 주로 절제 실험 표다. 중요도 맵 구성 요소를 비교한 표에서 CFG 기반 방식은 CLIP-T 30.8801, SSIM 0.7137, PSNR 17.8157, FID 44.6070이었고, 시간 축만 쓰면 31.0355/0.8109/20.5021/42.5269, 공간 축만 쓰면 31.1318/0.8110/20.5029/42.5736, 두 축을 모두 쓰면 31.1375/0.8112/20.5136/42.4474로 가장 좋았다. 시간 신호를 원시 score slope로 바꾸면 31.1166/0.8028/20.3944/42.9554로 모든 지표가 나빠지고 시각적 아티팩트가 생겼다. 제약 토큰 선택을 바꾼 절제에서는 시간↑∩공간↑ 조합이 SSIM 0.8112, LPIPS 0.1501, FID 42.4474로 최선이었고, 공간 응답이 낮은 토큰을 제약으로 쓰면 0.8005/0.1641/43.8062로 성능이 떨어졌다. 압축률 0.3에서 20개 PixArt-α 사례로 32명이 참여한 선호도 조사에서도 DaTMe가 텍스트 정렬과 시각 품질 모두 최고 평가를 받았다. 다만 PixArt-α와 FLUX의 주 실험 표에 담긴 구체적 수치는 제공된 본문 범위에 나오지 않는다.

개발자 관점에서 이 논문은 가속 기법의 병목이 연산량이 아니라 중요도 판단 신호라는 점을 보여준다. 캐싱이나 프루닝과 달리 토큰 병합은 어떤 토큰을 희생할지 결정하는 규칙 하나로 화질이 갈리므로, 프롬프트 정렬 점수만으로 토큰을 고르면 프롬프트에 없는 배경 구조가 먼저 무너진다는 것을 실무에서 확인할 수 있다. Tweedie 추정값은 이미 샘플러 내부에서 계산되는 값이라 추가 학습 없이 얻을 수 있고, CFG를 쓰지 않는 FLUX 계열에서도 동작한다는 점이 실용적이다. 다만 s=4, P=4 같은 하이퍼파라미터와 어느 레이어에 적용할지가 결과에 영향을 주므로, 자체 파이프라인에 붙일 때는 압축률별 SSIM·LPIPS를 먼저 재확인하는 편이 좋다.

저자들이 밝힌 한계로는 DaTMe가 지연 시간 면에서 기존 방법보다 약간 느리다는 점이 있다. 대신 압축률이 올라갈수록 기존 방법들이 FID를 비롯한 품질에서 크게 무너지는 반면 DaTMe는 상대적으로 견고해 전체적인 품질-효율 절충이 더 낫다고 주장한다. 또한 제약 집합을 두 축의 교집합으로 정의하는 방식은 엄격한 교집합이라 소스 집합이 목표 예산에 못 미치는 경우가 생기고, 이를 격자 샘플링으로 메우는 보완 절차가 필요하다는 구조적 전제를 갖는다.