Diffusion Templates: 제어 가능한 확산 모델을 위한 통합 플러그인 프레임워크

기존의 이미지 생성 제어 기술들은 특정 모델에 종속되어 서로 섞어 쓰거나 재사용하기 어려웠다. 이 논문은 제어 기능을 독립적인 '템플릿'으로 분리하여, 마치 레고 블록을 조립하듯 다양한 제어 기능을 하나의 모델에 자유롭게 추가하고 결합할 수 있는 표준 체계를 제시한다.

Diffusion 모델의 생성 과정은 노이즈에서 이미지를 찾아가는 반복적인 연산이며, 이 과정에서 Attention 메커니즘이 이미지의 구조와 특징을 결정한다. 기존의 ControlNet 등은 모델 옆에 별도의 신경망을 붙여 매 단계마다 간섭하는 방식을 취했는데, 이는 모델 구조가 바뀌면 재설계해야 하는 한계가 있었다.

Diffusion Templates는 이 간섭의 결과물을 KV-Cache라는 표준화된 데이터 형태로 변환한다. Transformer 기반 모델이 다음 단계의 연산을 수행할 때, 미리 계산된 Template Cache를 입력값 옆에 이어 붙임으로써 모델은 자연스럽게 외부에서 주입된 제어 정보를 자신의 연산 일부로 받아들이게 된다.

이 방식은 모델의 핵심 로직을 건드리지 않으면서도 강력한 제어를 가능하게 한다. 특히 여러 개의 템플릿을 동시에 사용할 때 각 템플릿이 생성한 캐시를 시퀀스 차원에서 결합하기만 하면 되므로, 복합적인 제어 기능을 구현하는 비용이 획기적으로 줄어든다.

프레임워크는 Template model, Template cache, Template pipeline의 세 가지 핵심 요소로 구성된다. Template model은 특정 작업의 입력(예: 깊이 지도, 스칼라 값)을 받아 중간 표현형인 Template cache로 변환하는 역할을 수행한다.

Template cache는 시스템 수준의 인터페이스로 정의된다. 주로 KV-Cache 형식을 권장하며, 이는 [입력 시퀀스 → Attention 레이어의 Key/Value 행렬 생성 → 기존 캐시와 결합] 과정을 거쳐 모델의 생성 방향을 결정한다. LoRA 파라미터 역시 이 인터페이스를 통해 전달될 수 있다.

Template pipeline은 추론 시 하나 이상의 템플릿을 로드하고 병합한다. 다중 제어 시 각 템플릿에서 나온 캐시들을 시퀀스 차원에서 연결(Concatenation)하여 기반 모델에 주입한다. 이 과정은 모델의 반복적인 Denoising 루프 외부에서 실행되므로 추가적인 연산 오버헤드가 매우 적으며 효율적인 추론이 가능하다.

FLUX.2-klein-base-4B 모델을 기반으로 다양한 실험을 진행했다. 이미지 편집 템플릿의 경우, 기반 모델의 고유 편집 기능과 대등한 품질을 유지하면서도 시퀀스 길이를 최적화하여 약 1.8배의 추론 속도 향상을 달성했다.

밝기 및 색상 조절 실험에서는 단순한 RGB 스케일링이 아닌 신경망 기반의 소프트 제어를 통해 시각적으로 자연스러운 조명 변화를 구현했다. 특히 미적 정렬(Aesthetic Alignment) 템플릿은 단 90쌍의 이미지 데이터셋 학습만으로도 학습 범위를 넘어서는 미적 개선 효과를 보여주며 일반화 성능을 입증했다.

다중 템플릿 융합(Fusion) 실험에서는 구조 제어, 이미지 편집, 색상 조절을 동시에 적용했을 때 각 기능이 충돌 없이 조화롭게 작동함을 확인했다. 이는 복잡한 예술적 스타일 변환과 구조적 제어가 단일 파이프라인에서 가능함을 시사한다.

Diffusion Templates 아키텍처는 기반 모델의 가중치를 동결(Frozen)한 상태에서 학습 가능한 사이드 브랜치를 부착하는 방식을 채택한다. 학습 목적 함수는 기반 모델의 원래 Pre-training Loss와 동일하게 설정하여 학습 목표를 일관되게 유지한다.

구현 측면에서 Stage I(입력 처리)과 Stage II(그래디언트 전파)를 분리했다. Stage I에서는 입력 데이터를 처리하여 재사용 가능한 중간 특징을 생성하고 이를 캐싱함으로써 반복적인 연산을 제거하고 학습 효율을 높였다.

KV-Cache 기반의 융합 메커니즘은 시퀀스 차원의 결합을 통해 이루어진다. 이는 모델의 Denoising 루프 내부가 아닌 외부에서 캐시를 준비하여 주입하는 방식이므로, 템플릿 개수가 늘어나도 GPU 메모리 점유율 상승을 억제할 수 있는 On-demand 로딩 전략을 지원한다.

현재 프레임워크는 프로토타입 단계이며, KV-Cache와 LoRA 외에 더 효율적인 인터페이스 포맷에 대한 탐구가 필요하다. 또한 현재는 정적 이미지 생성에 집중되어 있어 비디오 생성 모델에서의 시간적 일관성 제어에 대한 검증이 향후 과제로 남아있다.