기술 요약: 확산 모델 해석을 위한 잔차화 시간적 희소 오토인코더

문제 정의

텍스트-이미지 확산 모델은 반복적인 노이즈 제거 과정을 통해 이미지를 생성하므로, 내부 신경 계층은 단일 정적 표현이 아닌 활성화 궤적을 생성합니다. 희소 오토인코더 (SAE) 가 최근 확산 활성화를 해석 가능한 특징 방향으로 분해하는 데 적용되어 왔지만, 기존 접근법은 일반적으로 고립된 시간 단계에서의 활성화를 분석하거나, 시간을 조건으로 하거나, Diffusion Transformer 와 같은 아키텍처에 SAE 를 적용합니다.

이러한 기존 방법들의 결정적인 한계는 확산 활성화에 내재된 강한 시간적 상관관계를 명시적으로 고려하지 않는다는 점입니다. 인접한 시간 단계의 활성화는 매우 높은 상관관계를 가지며, 한 단계의 신호 대부분은 이전 단계에서 선형적으로 예측 가능합니다. 원시적이고 시간적으로 연결된 활성화에 직접 SAE 를 학습시키는 것은 모델 용량을 특징의 근본적인 희소 구조를 포착하는 대신 이러한 중복적이고 선형적으로 예측 가능한 동역학을 학습하는 데 낭비할 위험이 있습니다.

방법론

저자들은 선형적으로 예측 가능한 동역학을 잔차 구조와 명시적으로 분리함으로써 활성화 궤적에서 희소 특징을 학습하도록 설계된 잔차화 시간적 SAE를 제안합니다.

1. 데이터 수집 및 전처리

• 출처: 활성화는 classifier-free guidance 중 Stable Diffusion 1.5 의 U-Net 조건부 분기에서 수집됩니다. 구체적으로, mid-block cross-attention 계층의 잔차 기여도가 추출됩니다.
• 궤적 구성: 생성된 이미지의 경우, 활성화는 denoising 단계마다 $\ell$간격으로 하위 샘플링되어 활성화 텐서 시퀀스 $\{a(\tau_i)\}$를 형성합니다 (예: $\ell=10$).

2. 선형 잔차화

시간적 중복성을 제거하기 위해, 이 방법은 연속된 정규화된 활성화 블록 간에 선형 예측기를 적합시킵니다:

• 정규화: 활성화는 학습 세트 통계를 사용하여 시간 단계별로 정규화됩니다.
• 릿지 회귀: 각 인접 전이 $i-1 \to i$에 대해, 릿지 회귀기 $(W_i, b_i)$가 이전 활성화 $\bar{a}_{q,i-1}$로부터 현재 활성화 $\bar{a}_{q,i}$를 예측하도록 적합됩니다.
• 잔차 계산: 잔차 $r_{q,i} = \bar{a}_{q,i} - (W_i \bar{a}_{q,i-1} + b_i)$는 이전 단계에서 선형적으로 예측할 수 없는 활성화 성분을 포착합니다.

3. 잔차화 궤적에 대한 SAE 학습

SAE 의 입력은 초기 활성화와 후속 잔차로 구성된 연결된 벡터 $z_q$입니다:
$$z_q = [\bar{a}_{q,0}, r_{q,1}, \dots, r_{q,T-1}]$$
이 잔차화된 표현에 BatchTopK SAE가 학습됩니다. 목적은 희소성을 강제 (배치당 상위 $K$개의 활성 잠재 변수만 유지) 하면서 정규화된 입력을 재구성하는 것입니다. 이는 희소 잠재 변수가 선형적으로 예측 가능한 시간적 드리프트를 넘어서 존재하는 구조를 포착하도록 장려합니다.

4. 활성화 공간 해석

이 방법의 핵심 구성 요소는 학습된 특징을 원래 활성화 공간으로 매핑하여 궤적으로 분석하는 것입니다:

1. 디코더 방향: SAE 디코더 벡터는 초기에 정규화된 잔차화 공간에 있습니다.
2. 정규화 해제: 스케일이 적용되어 정규화된 잔차 공간으로 복귀합니다.
3. 선형 동역학 통합: 잔차 방향은 적합된 선형 예측기 ($W_i$) 를 통해 다시 매핑되어 정규화된 활성화 공간에서 전체 궤적 방향을 재구성합니다.
4. 최종 매핑: 결과는 원래 활성화 단위의 디코더 궤적 $\{\phi_{k,i}\}$를 얻기 위해 정규화 해제됩니다. 이를 통해 각 잠재 변수를 denoising 시간에 따라 진화하는 특징 궤적으로 분석할 수 있습니다.

주요 기여

1. 잔차화 시간적 SAE: 확산 활성화 궤적을 선형적으로 예측 가능한 성분과 잔차로 분리하고, 중복된 시간적 동역학을 모델링하지 않도록 후자에 대해 구체적으로 SAE 를 학습하는 방법론의 도입.
2. 활성화 공간 디코더 해석: 잔차화된 디코더 방향을 원래 활성화 공간으로 매핑하여 각 잠재 변수를 시공간적 특징 궤적으로 분석할 수 있게 하는 유도.
3. 경험적 특성화: Stable Diffusion 1.5 에 대한 포괄적인 평가를 통해 이러한 표현이 시간적으로 구조화되고 공간적으로 국소화된 방향을 포착하며 의미 있는 생성 효과를 가짐을 입증.

결과 및 평가

재구성 성능

• 잔차화의 우위성: 고정된 궤적 수준 희소성 예산 하에서, 잔차화 변형은 평균 제곱 오차 (MSE) 와 설명 분산 (EV) 측면에서 비잔차화 대응물 (연결형 및 시간 단계별 모두) 보다 일관되게 우월한 성능을 보입니다.
• 효율성: 잔차화 연결형 모델은 전체 궤적에 대해 단일 희소 코드를 사용하면서도 경쟁력 있는 재구성 성능을 달성하는 반면, 시간 단계별 기준선은 분산된 코드가 필요합니다.
• 기준선 비교: 이 방법은 재구성 품질 측면에서 Matryoshka SAE 기준선 (비잔차화, 성분 정규화 궤적을 사용) 보다 우월합니다.

시공간 구조 분석

• 시간적 전문화: 잠재 변수는 뚜렷한 시간적 프로파일을 보입니다. "초기" 특징은 광범위한 시간적 지지와 높은 시간 간격 자기 유사성을 가지며, 이는 지속적인 구조적 개념을 포착함을 시사합니다. "후기" 특징은 더 많은 시간 단계 특이성을 가지며 시간적으로 덜 안정적입니다.
• 공간적 국소화:
  • 초기 특징은 더 넓은 양의 공간적 지지 (높은 공간 엔트로피) 를 보이며, 종종 확산되거나 경계와 같은 영역을 덮습니다.
  • 중간 및 후기 특징은 더 공간적으로 집중되어 (낮은 엔트로피) 특정 객체와 같은 영역이나 중요한 영역에 국소화됩니다.

조향 실험

• 단일 특징 조향: 특정 잠재 변수를 개입시키면 가시적인 생성 효과가 발생합니다. 초기 특징은 글로벌 구조에 영향을 미치는 반면, 중간 특징은 변화를 국소화할 수 있습니다 (예: 몸통에 보석 추가). 후기 특징은 일부 예시에서 덜 명확한 의미적 효과를 보였습니다.
• 특징 전이: RIEBench 데이터셋을 사용하여 저자들은 소스-대상 특징 전이를 수행했습니다. 잔차화 연결형 모델은 편집 효율성 (대상 프롬프트에 대한 CLIP 유사성 향상) 과 소스 이미지 구조 보존 (LPIPS 거리 최소화) 사이의 가장 좋은 균형을 달성했습니다. 이는 편집 효율성 측면에서 비잔차화 및 시간 단계별 변형보다 우월했습니다.

중요성 및 주장

이 논문은 잔차화 시간적 SAE가 시간적으로 구조화된 확산 활성화를 연구하는 데 유용한 프레임워크를 제공한다고 주장합니다. 선형적으로 예측 가능한 드리프트를 명시적으로 제거함으로써, 이 방법은 희소 잠재 변수가 생성 과정의 비중복적이고 의미 있는 구조에 집중할 수 있게 합니다.

저자들은 그들의 접근 방식이 고립된 시간 단계나 비잔차화 시간적 표현이 아닌 잔차화 활성화 궤적에 대해 희소 사전 (dictionary) 을 학습한다는 점에서 기존 작업과 다르다고 강조합니다. 이러한 잠재 변수를 궤적으로 활성화 공간으로 다시 매핑할 수 있는 능력은 denoising 과정에서 특정 특징이 언제 그리고 어디에서 발현되는지에 대한 더 깊은 이해를 가능하게 합니다.

저자가 지적한 한계:

• 실험은 Stable Diffusion 1.5 의 단일 U-Net 활성화 소스 (mid-block cross-attention) 에 초점을 맞추고 있습니다.
• 조향 실험은 정성적입니다. 논문은 최적화된 이미지 편집 성능을 주장하는 것이 아니라, 학습된 특징의 해석 가능성과 인과적 영향을 검증하기 위해 조향을 사용합니다.