Published on

확산 모델 추론 확장 새로운 패러다임

저자
  • avatar
    이름
    Ajax
    Twitter

확산 모델 추론 확장: 새로운 패러다임

최근 대규모 언어 모델(LLMs)의 발전은 추론 과정에서 확장의 효과를 입증했습니다. o1, o3, DeepSeek R1, QwQ, Step Reasoner mini와 같은 모델들은 추론 중 계산량 증가가 성능을 크게 향상시킬 수 있음을 보여주었습니다. 이는 확산 모델에도 이러한 원칙을 적용할 수 있는지에 대한 의문을 제기합니다.

뉴욕대학교의 Xie Saining 팀은 이러한 질문을 탐구했습니다. 일반적인 검색 프레임워크를 사용한 체계적인 조사를 통해, 추론 시간 확장이 확산 모델에도 효과적임을 발견했습니다. 추론 중 계산량 증가는 생성된 샘플 품질을 크게 향상시킵니다. 또한, 이미지의 복잡성으로 인해 다양한 응용 시나리오에 맞게 프레임워크 내에서 구성 요소의 다양한 조합이 가능합니다.

주요 연구 결과

  • 추론 시간 확장은 확산 모델에 효과적입니다: 추론 중 더 많은 계산 자원을 할당하면 더 높은 품질의 샘플을 얻을 수 있습니다.
  • 구성 요소 조합의 유연성: 프레임워크는 다양한 응용 분야에 맞춰 다양한 구성 요소 구성을 허용합니다.
  • 디노이징 단계 그 이상: 연구 결과에 따르면, 샘플링 중 더 나은 노이즈를 검색하는 것이 단순히 디노이징 단계를 늘리는 것 이상으로 NFE를 확장하는 또 다른 차원입니다.
  • 두 가지 설계 축: 프레임워크는 두 가지 주요 설계 축에 초점을 맞춥니다.
    • 검증자(Verifiers): 검색 과정 중 피드백을 제공합니다.
    • 알고리즘(Algorithms): 더 나은 노이즈 후보를 찾습니다.

연구 방법론

연구팀은 다양한 사용 사례를 시뮬레이션하기 위해 세 가지 시나리오에서 검증자를 탐색했습니다.

  1. 최종 평가에 대한 특권 정보가 제공되는 시나리오
  2. 생성을 안내하기 위해 조건부 정보가 제공되는 시나리오
  3. 추가 정보가 없는 시나리오

알고리즘의 경우, 다음을 조사했습니다.

  1. 무작위 검색(Random Search): 고정된 후보 집합에서 최적의 항목을 선택합니다.
  2. 제로차 검색(Zero-Order Search): 검증자 피드백을 사용하여 노이즈 후보를 반복적으로 개선합니다.
  3. 경로 검색(Path Search): 검증자 피드백을 사용하여 확산 샘플링 궤적을 반복적으로 개선합니다.

이 연구는 먼저 상대적으로 간단한 ImageNet 클래스 조건부 생성 설정에서 이러한 설계를 탐색했습니다. 그 후, 이러한 설계를 더 큰 규모의 텍스트 조건부 생성에 적용하고 제안된 프레임워크를 평가했습니다.

추론 시간 확장

이 논문은 확산 모델에서 추론 시간을 확장하기 위한 프레임워크를 제안하며, 최적의 샘플링 노이즈를 찾는 문제로 프레임워크를 구성합니다. 이 프로세스에는 두 가지 핵심 구성 요소가 포함됩니다.

  • 검증자(Verifiers): 생성된 샘플의 품질을 평가하는 사전 학습된 모델입니다. 생성된 샘플과 선택적 조건을 입력받아 스칼라 점수를 출력합니다.
  • 알고리즘(Algorithms): 검증자 점수를 사용하여 더 나은 후보 샘플을 찾는 알고리즘입니다. 검증자(V), 사전 학습된 확산 모델(_θ) 및 생성된 샘플과 조건 집합을 입력받아 최상의 초기 노이즈를 출력합니다.

총 추론 예산은 디노이징 단계와 검색 비용을 포함한 총 함수 평가 횟수(NFE)로 측정됩니다.

검색 검증자

연구자들은 선택된 샘플의 최종 평가에 대한 완전한 정보를 가진 오라클 검증자로 시작했습니다. ImageNet의 경우 FID 및 IS와 같은 메트릭이 포함되었습니다. 그런 다음 CLIP 및 DINO와 같은 더 접근하기 쉬운 사전 학습된 모델을 지도 검증자로 탐색했습니다. 이러한 모델은 샘플을 분류하는 데 사용되었으며, 클래스 레이블에 해당하는 가장 높은 로짓을 가진 샘플을 선택했습니다.

그러나 이러한 분류기는 점별로 작동하며 FID 점수의 목표와 부분적으로만 일치한다는 것을 발견했습니다. 이로 인해 계산이 증가함에 따라 샘플 분산이 감소하고 모드 붕괴가 발생했습니다. "검증자 해킹"이라고 불리는 이 현상은 무작위 검색 알고리즘의 제약 없는 검색 공간에 의해 가속화되었습니다.

흥미롭게도, 이 연구는 검증자가 검색을 효과적으로 안내하기 위해 반드시 조건부 정보가 필요하지 않다는 것을 발견했습니다. DINO/CLIP 분류기의 로짓과 낮은 노이즈 수준에서의 x 예측과 최종 깨끗한 샘플 간의 특징 공간의 코사인 유사성 사이의 강한 상관 관계를 관찰했습니다. 이를 통해 추가 조건부 정보가 필요하지 않으면서도 효과적인 확장 동작을 보이는 자기 지도 검증자를 사용하게 되었습니다.

검색 알고리즘

검증자 해킹을 완화하기 위해 연구자들은 후보 샘플을 점진적으로 최적화하는 더 세련된 검색 알고리즘을 탐색했습니다. 여기에는 제로차 검색 방법이 포함되었습니다.

  1. 무작위 가우시안 노이즈를 피벗 포인트로 시작합니다.
  2. 피벗 포인트의 주변에서 N개의 후보를 찾습니다.
  3. ODE 솔버를 통해 후보를 실행하여 샘플과 검증자 점수를 얻습니다.
  4. 최상의 후보로 피벗 포인트를 업데이트하고 1~3단계를 반복합니다.

또한 샘플링 궤적을 따라 검색할 가능성을 탐색하는 경로 검색 알고리즘을 조사했습니다.

  1. N개의 초기 노이즈 샘플을 샘플링하고 ODE 솔버를 노이즈 수준 σ까지 실행합니다.
  2. 각 샘플에 노이즈를 추가하고 정방향 노이징 프로세스를 시뮬레이션합니다.
  3. 각 노이즈 샘플에서 ODE 솔버를 실행하고 검증자 점수를 기준으로 상위 N개 후보를 유지하고 ODE 솔버가 σ=0에 도달할 때까지 반복합니다.
  4. 나머지 N개 샘플을 무작위로 검색하고 최상의 샘플을 유지합니다.

제로차 및 경로 검색 알고리즘은 모두 무작위 검색에 비해 강력한 지역성을 유지합니다.

텍스트-이미지 시나리오에서 확장

연구팀은 대규모 텍스트-이미지 작업에서 검색 프레임워크의 확장 기능을 조사했습니다. 평가를 위해 DrawBench 및 T2I-CompBench 데이터 세트를 사용했으며, FLUX.1-dev 모델을 백본으로 사용했습니다. 또한 Aesthetic Score Predictor, CLIPScore 및 ImageReward를 포함하여 지도 검증자 선택을 확장했습니다. 또한 이러한 세 가지 검증자를 결합하여 검증자 앙상블을 만들었습니다.

분석: 검증자-작업 정렬

이 연구는 다양한 데이터 세트에서 다양한 검증자-알고리즘 조합의 결과를 비교했습니다. DrawBench에서 모든 검증자를 사용하는 것이 일반적으로 샘플 품질을 향상시킨다는 것을 발견했습니다. 그러나 Aesthetic 및 CLIP 검증자를 개별적으로 사용하면 편향에 과적합되어 서로에게 부정적인 영향을 미칠 수 있다는 점을 관찰했습니다. 이는 평가 초점의 불일치에서 비롯됩니다. Aesthetic Score는 시각적 품질에 초점을 맞춰 고도로 양식화된 이미지를 선호하는 반면, CLIP은 시각적-텍스트 정렬을 우선시하여 때로는 시각적 품질을 희생합니다. 일부 검증자는 특정 작업에 더 적합하며, 검증자의 효과는 작업 요구 사항과의 정렬에 따라 달라진다는 점을 지적했습니다.

알고리즘 성능

세 가지 검색 알고리즘(무작위, 제로차 및 경로)은 모두 DrawBench에서 샘플링 품질을 효과적으로 개선했습니다. 그러나 무작위 검색은 다른 두 방법의 지역적 특성으로 인해 일부 측면에서 더 나은 성능을 보였습니다. 무작위 검색은 검증자 편향에 더 빠르게 수렴한 반면, 다른 두 알고리즘은 최적이 아닌 후보에 대한 개선이 필요합니다.

미세 조정과의 호환성

연구팀은 미세 조정된 모델과의 검색 방법의 호환성을 조사했습니다. DPO 미세 조정된 Stable Diffusion XL 모델을 사용했으며, 검색 방법이 다른 모델로 일반화될 수 있으며 이미 정렬된 모델의 성능을 향상시킬 수 있음을 발견했습니다.

추론 계산의 다양한 차원의 효과

이 연구는 추론 계산의 다양한 측면이 결과에 어떤 영향을 미치는지 탐구했습니다.

  • 검색 반복 횟수: 반복 횟수를 늘리면 노이즈가 최적 값에 더 가까워집니다.
  • 검색 반복당 계산: 반복당 디노이징 단계 수를 조정하면 계산적으로 최적인 영역이 다르게 나타납니다.
  • 최종 생성 계산: 연구팀은 최종 디노이징 단계에 최적 설정을 사용하여 최종 샘플 품질이 가장 높도록 보장했습니다.

계산 투자 효과

연구자들은 더 작은 확산 모델에서 추론 시간 확장의 효과를 탐구했습니다. ImageNet의 경우 더 작은 모델을 확장하는 것이 매우 효율적일 수 있다는 것을 발견했습니다. 경우에 따라 더 작은 모델에서 검색하는 것이 검색 없이 더 큰 모델보다 성능이 뛰어날 수 있습니다. 그러나 효과는 더 작은 모델의 기준 성능에 따라 달라집니다.

텍스트 기반 설정에서 PixArt-Σ는 계산량의 일부만 사용하여 FLUX-1.dev보다 성능이 뛰어났습니다. 이러한 결과는 훈련 중에 소비되는 상당한 계산 자원을 생성 중에 더 적은 양의 계산으로 상쇄하여 더 높은 품질의 샘플을 더 효율적으로 생성할 수 있음을 보여줍니다.