Papers·5일 전
Tsinghua-LeapLab, Softmax attention 을 TTT 로 변환해 SD3.5 추론 1.32~1.47x 가속
Tsinghua-LeapLab 팀이 Softmax attention 을 linear-complexity TTT(Test-Time Training) 구조로 변환하는 방법을 제안했습니다. 두 계층 동적 구조가 Softmax attention 과 구조적으로 정렬된다는 점을 활용해 사전 학습된 가중치를 직접 계승하고, key instance normalization 과 locality enhancement 모듈로 표현 정렬을 추가했습니다. Stable Diffusion 3.5 에 적용한 결과, 4×H20 GPU 에서 1시간 fine-tuning 만으로 원본과 유사한 품질을 유지하면서 1K/2K 해상도에서 각각 1.32배, 1.47배 추론 속도 향상을 달성했습니다. 코드는 공개되었습니다.
Tsinghua-LeapLab 이 Softmax attention 을 linear-complexity TTT 구조로 변환해 SD3.5 추론을 1.32~1.47배 가속했습니다.
핵심 결론
- 태스크 — Stable Diffusion 3.5 의 Softmax attention 을 linear-complexity TTT 로 변환.
- 성능 — 4×H20 GPU, 1시간 fine-tuning 으로 원본 대비 동등한 품질, 1K/2K 해상도에서 각각 1.32배, 1.47배 추론 속도 향상.
방법
- 구조 정렬 — TTT 의 두 계층 동적 formulation 이 Softmax attention 과 구조적으로 정렬되어 사전 학습 가중치를 직접 계승 가능.
- 표현 정렬 — key instance normalization 으로 key shift-invariance 를, lightweight locality enhancement module 로 locality 를 맞춤.
- 변환 대상 — SD3.5 의 attention layer 를 TTT layer 로 교체 후 fine-tuning.
한계·조건
- 리소스 — 4×H20 GPU, 1시간 fine-tuning 으로 충분하지만, 사전 학습된 Transformer 가 있어야 함.
- 범위 — SD3.5 에 한정된 검증이며, 다른 diffusion 모델이나 LLM 으로의 일반화는 추가 실험 필요.
- 코드 — GitHub 공개 — https://github.com/LeapLabTHU/Transformer-to-TTT
편집자 한 줄
사전 학습된 가중치를 버리지 않고 linear attention 으로 전환할 수 있다는 점이 실용적입니다. 다만 SD3.5 외 도메인에서도 같은 추세가 유지될지는 지켜볼 만합니다.
- #attention
- #linear-complexity
- #ttt
- #stable-diffusion
- #tsinghua
Tsinghua-LeapLab