기술 요약: Sparse Context – 토큰 드롭핑을 통한 효율적인 참조 조건부 생성

문제 정의

참조 기반 확산 모델(Reference-based diffusion models)은 입력 이미지에 조건을 부여하여 이미지 편집 및 피사체 중심 합성(subject-driven synthesis)과 같은 작업을 가능하게 함으로써 제어 가능한 이미지 생성 분야에서 핵심적인 역할을 해왔습니다. 그러나 이러한 모델들은 심각한 계산 병목 현상에 직면해 있습니다. 현대적인 트랜스포머 기반 확산 모델(DiTs)에서 참조 이미지는 일반적으로 VAE를 통해 밀집된 공간 그리드 형태의 토큰으로 인코딩된 후 노이즈가 섞인 입력 토큰과 결합됩니다. 셀프 어텐션(self-attention) 연산은 토큰 수에 따라 제곱 비례하여 확장되기 때문에, 하나 이상의 고해상도 참조 이미지를 포함하는 것은 추론 시간과 메모리 사용량을 급격히 증가시킵니다. 이러한 비용은 여러 개의 참조 이미지를 사용하는 실제 응용 환경에서 감당하기 어려운 수준이 됩니다. 프롬프트 기반 생성(prompt-driven generation)을 위한 효율화 기술은 존재하지만, 시각적 입력의 밀집된 표현 특성으로 인해 독특한 과제를 안고 있는 참조 조건부 모델에 대해서는 아직 충분히 탐구되지 않았습니다.

방법론

본 논문은 참조 토큰의 축소된 하위 집합만을 유지함으로써 희소한 참조 표현을 구축하도록 설계된 Sparse Context 프레임워크를 소개합니다. 이 방법론은 세 가지 핵심 구성 요소를 통해 문제를 해결합니다.

1. 중복성에 대한 관찰

저자들은 먼저 참조 토큰이 상당한 중복성을 포함하고 있음을 입증했습니다. 사전 학습된 FLUX.2-Klien 모델에 대해 "나이브 컨텍스트 프루닝(Naïve Context Pruning)" 실험을 수행한 결과, 추론 시 참조 토큰의 최대 80%를 무작위로 드롭하더라도 미세한 디테일의 손실은 발생하지만 장면의 대략적인 레이아웃은 대체로 유지됨을 보여주었습니다. 이는 참조 조건부 모델이 본질적으로 희소한 입력에 대해 강건(robust)하다는 것을 시사하지만, 전체 참조와 부분 참조 사이의 성능 격차는 훈련-테스트 간의 분포 불일치(distribution mismatch)로 인해 악화된다는 점을 보여줍니다.

2. 훈련 전략: 무작위 토큰 드롭핑 (Random Token Dropping)

이러한 분포 격차를 해소하고 강건한 희소 생성을 가능하게 하기 위해, 저자들은 무작위 토큰 드롭핑을 사용하여 최첨단 참조 조건부 생성기를 미세 조정(fine-tuning)합니다.

• 과정: 훈련 중에 참조 토큰의 무작위 비율이 드롭됩니다 (토큰 유지 비율 $f \in [0.05, 0.25]$). 이 토큰들은 노이즈가 섞인 타겟 토큰과 결합되기 전에 드롭됩니다.
• 목표: 이를 통해 모델이 부분적인 참조 정보로도 효과적으로 작동하는 강건한 표현을 학습하도록 유도하며, 특정 토큰 선택 규칙으로부터 모델을 분리합니다.
• 구현: 지시어 기반 편집 및 개인화(personalization)를 다루는 큐레이션된 데이터셋을 사용하여 FLUX-2-Klein-9B 베이스 모델에 경량 LoRA(Low-Rank Adaptation)를 적용합니다.

3. 추론 전략: 태스크 인지형 토큰 선택 (Task-Aware Token Selection)

훈련에서는 강건성을 확보하기 위해 무작위 드롭을 사용하지만, 추론 시에는 고정된 계산 예산 내에서 품질을 극대화하기 위해 태스크 인지형 토큰 선택 전략을 사용합니다.

• 지시어 기반 이미지 편집: 공간 구조를 보존하기 위해 에지 인식 샘플링(edge-aware sampling) 방식을 사용합니다. 참조 이미지에서 Canny 에지 맵을 추출하고, 균일한 영역보다 구조적 경계(structural boundaries)를 우선시하는 확률 분포를 기반으로 토큰을 샘플링합니다.
• 개인화: 피사체의 정체성을 보존하기 위해 **살리언시 기반 샘플링(saliency-based sampling)**을 사용합니다. 살리언시 맵이 피사체 마스크를 근사하여, 모델이 참조 이미지의 가장 정보량이 많은 영역에 토큰을 집중할 수 있도록 합니다.

주요 기여

1. Sparse Context 프레임워크: 기본 모델의 아키텍처를 수정하지 않고도 토큰을 드롭함으로써 희소한 참조 표현을 구축하는 새로운 접근 방식을 제시하여, 셀프 어텐션의 이차적(quadratic) 비용을 크게 줄입니다.
2. 강건한 미세 조정 전략: 무작위 토큰 드롭핑을 사용하는 훈련 체계를 통해 모델이 다양한 수준의 희소성에도 일반화될 수 있도록 하여, 추론 시 계산 예산을 유연하게 조절할 수 있게 합니다.
3. 태스크 인지형 추론: 무작위 드롭핑보다 태스크에 중요한 정보를 우선시하는 특정 토큰 선택 휴리스틱(편집을 위한 에지 기반, 개인화를 위한 살리언시 기반)을 제안합니다.
4. 기존 방법론과의 직교성: Sparse Context가 Token Merging(ToMe) 및 KV-캐싱과 같은 다른 효율화 기술과 호환되며 상호 보완적임을 입증하여, 결합된 속도 향상 효과를 얻을 수 있음을 보여줍니다.

실험 결과

본 방법론은 FLUX-2-Klein-9B 모델을 사용하여 지시어 기반 이미지 편집(PIE-Bench) 및 참조 기반 개인화(Subject-200K, CustomDiffusion-105)에 대해 평가되었습니다.

• 효율성:
  • 단일 참조: 최대 2배의 추론 속도 향상을 달s합니다 (예: 토큰 유지율 10%에서 1.83배).
  • 다중 참조: 참조 토큰이 어텐션 계산의 대부분을 차지하는 시나리오(예: 6개의 참조 이미지)에서 최대 4배의 추론 속도 향상을 달성합니다.
  • 결합된 효율성: 다중 참조 생성을 위해 KV-캐싱과 통합했을 때, 최대 6.72배의 속도 향상을 달성합니다.
• 품질:
  • 이미지 편집: Sparse Context는 나이브 컨텍스트 프루닝 베이스라인보다 장면 구조와 정체성을 훨씬 더 잘 보존합니다. 정량적 지표(LPIPS, DINO, CLIP-I)는 Sparse Context가 낮은 토큰 비율(예: $f=0.1$)에서도 높은 충실도를 유지하는 반면, 베이스라인은 정체성 훼손을 겪음을 보여줍니다.
  • 개인화: 살리언시 기반 샘플링은 무작위 드롭핑 시 손실되는 미세한 피사체 디테일(예: 특정 색상, 패턴)을 성공적으로 보존합니다.
• 통합: Sparse Context를 Token Merging과 결합하면 성능이 더욱 향상되며, 이는 미세 조정된 모델이 추가적인 토큰 감소 기술을 효과적으로 활용함을 입증합니다.

의의 및 주장

본 논문은 Sparse Context가 참조 기반 확산 모델에 내재된 미개척 효율성 잠재력을 끌어올린다고 주장합니다. 밀집된 토큰 그리드가 종종 불필요하다는 점을 인식함으로써, 저자들은 제어 가능성이나 시각적 충실도를 저해하지 않고도 참조 기반 생성을 훨씬 더 효율적으로 만들 수 있음을 입증했습니다.

저자들은 이 접근 방식이 매우 단순하고 모듈화되어 있으며, 경량 미세 조정만을 필요로 하고 추론 시 유연한 제어를 제공한다는 점을 강조합니다. 이 방법은 고해상도, 다중 참조 생성 모델의 배포에 있어 결정적인 병목 현상을 해결하며, 이를 통해 실제 응용 환경에서의 실용성을 높인다고 주장합니다. 본 연구는 참조 표현의 중복성이 활용 가능한 근본적인 속성임을 강조하며, 이러한 발견은 테스트된 특정 모델을 넘어 DiT 기반의 이미지 및 비디오 생성 시스템 전반으로 확장될 수 있습니다.