HybridStitch: 확산 모델 가속화를 위한 픽셀 및 타임스텝 수준의 모델 스티칭

대형 확산 모델은 고품질 이미지를 생성하지만 막대한 연산량으로 인해 실시간 응용에 한계가 있다. 이 논문은 이미지 내의 복잡한 영역은 큰 모델이, 단순한 영역은 작은 모델이 처리하도록 픽셀 단위로 역할을 분담시켜 품질 저하 없이 생성 속도를 획기적으로 높였다.

왜 중요한가

핵심 기여

영역 인식형 모델 스위칭 패러다임 제안

이미지 전체를 동일하게 처리하는 기존 방식과 달리 픽셀별 렌더링 난이도에 따라 대형 모델과 소형 모델을 동적으로 전환하는 HybridStitch 기법을 도입했다.

KV 캐시 재사용을 통한 컨텍스트 보존

대형 모델이 이미지의 일부인 마스크 영역만 처리할 때 발생하는 컨텍스트 손실 문제를 이전 단계의 Key-Value 캐시를 활용해 해결함으로써 생성 품질을 유지했다.

훈련이 필요 없는 가속화 프레임워크 구축

추가적인 모델 학습이나 미세 조정 없이 기존의 사전 학습된 대형 및 소형 모델 조합만으로 즉시 적용 가능한 Train-free 가속화 방식을 구현했다.

핵심 아이디어 이해하기

Diffusion 모델은 노이즈에서 이미지를 복원하는 과정을 거치는데 모든 픽셀이 매 단계마다 동일한 수준의 연산을 필요로 하지는 않는다. 배경처럼 단순한 부분은 일찍 완성되는 반면 사물의 세부 묘사가 필요한 부분은 더 많은 정제가 필요하다. 기존의 가속 기법들은 특정 단계 전체를 작은 모델로 대체했지만 이는 이미지 내의 지역적 차이를 무시한다는 한계가 있었다. HybridStitch는 이 문제를 해결하기 위해 하이브리드 단계를 도입한다. 먼저 대형 모델로 전체적인 구도를 잡은 뒤 두 모델의 예측값 차이가 큰 어려운 영역만 마스크로 지정한다. 이후 작은 모델은 전체 이미지를 빠르게 훑으며 초안을 만들고 대형 모델은 마스크된 영역만 집중적으로 정밀 묘사하는 방식으로 협업한다. 이 과정에서 대형 모델은 전체 이미지 대신 10~40%의 픽셀만 처리하므로 연산량이 급격히 줄어든다. 결과적으로 전체 이미지를 대형 모델로만 생성할 때보다 훨씬 빠르면서도 작은 모델만 썼을 때 발생하는 품질 저하를 효과적으로 방지한다.

방법론

HybridStitch는 세 단계로 구성된다. 첫 번째 단계는 대형 모델만을 사용하여 이미지의 기본 레이아웃을 구성한다. 두 번째 단계인 하이브리드 단계에서는 소형 모델이 전체 잠재 상태를 처리하여 초안을 만들고 대형 모델은 이전 단계의 예측 차이가 큰 영역인 마스크 영역만 선택적으로 정제한다. 마지막 단계에서는 소형 모델만을 사용하여 최종 이미지를 완성한다. 마스크 생성 및 업데이트는 인접한 두 타임스텝 사이의 잠재 변수 차이인 $D_t = \text{Mean}(\frac{||X_t - X_{t+1}||_1}{||X_t||_1})$ 를 계산하여 수행된다. [현재 단계 $X_t$ 와 다음 단계 $X_{t+1}$ 의 잠재 변수 값을 입력으로] → [두 값의 차이에 대한 L1 노름을 구하고 이를 $X_t$ 의 크기로 나누어 정규화한 뒤 전체 픽셀에 대해 평균을 내는 연산을 수행해] → [0에서 1 사이의 변화율 수치를 얻고] → [이 값이 설정된 임계값보다 낮아지면 다음 단계로 전환하거나 마스크 영역을 갱신하는 기준으로 활용한다.] 마스크된 영역만 처리할 때 발생하는 컨텍스트 부족 문제를 해결하기 위해 KV 캐시 기법을 적용한다. 대형 모델의 Attention 연산 시 현재 마스크된 토큰들로부터 Query를 생성하고 이를 이전 단계에서 저장된 전체 이미지의 Key, Value 캐시와 결합하여 연산한다. [마스크된 토큰과 이전 단계의 전체 KV 캐시를 입력으로] → [Attention 메커니즘을 통해 마스크 영역과 전체 컨텍스트 간의 관계를 계산해] → [정제된 특징 맵을 얻고] → [이를 통해 모델이 이미지의 일부만 보면서도 전체적인 맥락을 유지한 채 정교한 예측을 수행하게 한다.]

주요 결과

Stable Diffusion 3를 대상으로 한 실험에서 HybridStitch는 기존 SOTA 기법인 SRDiffusion 대비 최대 18.06%의 추가적인 지연 시간 감소를 달성했으며 전체적으로 1.83배의 속도 향상을 기록했다. COCO 데이터셋 기반 평가에서 FID 점수 30.34를 기록하여 T-Stitch(31.87)나 SRDiffusion(31.67)보다 우수한 이미지 품질을 입증했다. 다양한 마스크 비율에 대한 실험 결과 30% 마스크 설정에서 가장 최적의 효율성과 품질 균형을 보였다. 특히 고성능 GPU인 H100 및 A100 환경에서도 최소 1.5배 이상의 속도 향상을 유지하며 하드웨어 범용성을 확인했다.

기술 상세

HybridStitch 아키텍처는 대형 모델인 Stable Diffusion 3.5 Large와 소형 모델인 Stable Diffusion 3 Medium의 출력을 픽셀 단위로 결합하는 구조다. 하이브리드 단계에서 소형 모델의 출력값 중 마스크된 영역을 대형 모델의 출력값으로 교체하여 다음 단계의 입력으로 전달한다. 핵심 차별점은 영역 인식형 스위칭과 KV 캐시 기반 컨텍스트 보완이다. 기존 T-Stitch가 타임스텝 단위로만 모델을 교체하여 초기 단계의 오차가 누적되는 문제를 해결하기 위해 HybridStitch는 초기에는 대형 모델을 사용하고 이후 점진적으로 소형 모델의 비중을 높이되 복잡한 영역은 끝까지 대형 모델이 관리하도록 설계되었다. 구현 측면에서 마스크 업데이트는 매 단계마다 동적으로 이루어지며 이는 정적인 마스크를 사용할 때보다 FID 및 LPIPS 지표에서 유의미한 품질 향상을 가져온다. 또한 FFN이나 정규화 레이어는 토큰 간 상호작용이 없으므로 추가 연산 없이 마스크된 영역에만 적용하여 효율성을 극대화했다.

한계점

마스크 인덱싱 및 불연속적인 메모리 접근으로 인해 10% 마스크를 사용하더라도 실제 지연 시간이 이론적인 10% 수준까지 선형적으로 감소하지는 않는 오버헤드가 존재한다.

실무 활용

고해상도 이미지 생성이 필요한 서비스에서 서버 비용을 절감하고 사용자 대기 시간을 줄이는 데 즉시 적용 가능하다.

실시간 텍스트-이미지 생성 웹 서비스의 추론 가속화
모바일 기기 등 리소스 제한 환경에서의 대형 확산 모델 구동
대규모 이미지 데이터셋 생성을 위한 인프라 비용 최적화

코드 공개 여부: 비공개

키워드

Diffusion Model(확산 모델)Model Stitching(모델 스티칭)Inference Acceleration(추론 가속화)Region-aware Switching(영역 인식형 스위칭)KV Cache(KV 캐시)

용어 해설

디노이징 단계(Denoising Step): — Diffusion 모델이 노이즈에서 이미지를 생성하는 반복적인 단계이다. 각 단계마다 노이즈를 조금씩 제거하며 단계가 진행될수록 이미지의 세부 사항이 결정된다. 전체 추론 속도를 결정하는 핵심 요소다.
KV 캐시(KV Cache): — Transformer 구조에서 이전 단계의 Key와 Value 계산값을 저장해두는 기술이다. 매번 처음부터 계산하지 않고 저장된 값을 재사용하여 연산 속도를 높이고 컨텍스트를 유지하는 데 중요하다.
잠재 상태(Latent State): — 이미지를 직접 처리하는 대신 압축된 저차원 공간에서 표현한 상태이다. Diffusion 모델은 주로 이 잠재 공간에서 연산을 수행하여 메모리 효율성을 높이고 복잡한 패턴을 효과적으로 학습한다.
FID 점수(FID Score): — 생성된 이미지의 품질과 다양성을 평가하는 지표이다. 실제 이미지 분포와 생성된 이미지 분포 사이의 거리를 측정하며 값이 낮을수록 실제 이미지와 유사한 고품질임을 의미한다.
지각적 이미지 패치 유사도(LPIPS): — 두 이미지 사이의 시각적 유사도를 측정하는 지표이다. 단순 픽셀 차이가 아닌 딥러닝 모델의 특징 맵을 비교하여 인간이 느끼는 지각적 유사성을 수치화한다.

HybridStitch: 확산 모델 가속화를 위한 픽셀 및 타임스텝 수준의 모델 스티칭

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드