Abstract
최근 Backpropagation-based 접근법들은 denoising chain 안에서 reward function을 사용하는 diffusion model이 좀 더 잘 동작하는 것을 목표로한다. 그러나 이것은 상당한 계산 비용과 denoising chain의 길이로 인한 gradient explosion 위험 때문에, 존재하는 접근법들은 안좋은 결과를 내는 gradient backpropagation을 제대로 완료시키는 것에 고군분투하고있다. 이 논문에서, 더 짧은 denoising chain을 활용하여 효율적인 fine-tuning 전략을 사용하는 Shortcut-based Fine-Tuning (SHORTFT)을 도입한다.
Shortcut-based Fine-Tuning (SHORTFT)는 최근 연구된 trajectory-preserving few-step diffusion model을 사용하는데, 이것은 original denoising chain을 shortcut 할 수 있게 한다. 이 chain에서 최적화는 foundation model에서 fine-tuning의 효율성과 유효성을 향상시킨다. 논문의 방법은 철저하게 실험했고 다양한 보상함수에 대해 효과적으로 적용될 수 있다고 한다. 그리고 이 실험에서는 SOTA 모델을 능가하는 성능을 보인다고 말한다.
Introduction
Diffusion model들은 generative modeling을 위해 만들어졌고, 현실 사진처럼 text-to-image 합성과 같은 분야에 뛰어난 성능을 보여준다. 그러나 이 모델의 maximum likelihood training 목적함수는 종종 downstream task들에서 잘 동작하지 않는다고 한다.
- 작은 규모의, 사람이 선별한 데이터셋(예: LAION Aesthetics)에 대해 직접적인 감독 방식의 fine-tuning을 수행하는 것이 가장 단순한 해결책일 수 있지만, data collection의 비용과 데이터셋의 빠른 노후화로 인해, 이 접근 방법이 비실용적이라고 한다.
- LLM에서 Reinforcement Learning from Human Feedback (RLHF)의 성공적인 적용은 diffusion model에 reward function을 적용해 더 좋은 성능을 내도록 실험되어졌다. 이 기법은 특정 도메인에 대해 성능 향상을 보였지만, high-variance gradient와 여러 prompt에 대해 비효율적이고 제한적인 성능을 보여줬다.
- 최근 backpropagation-based 방법들은 reward function을 적용한 diffusion model이 더 좋은 성능을 내도록 했지만, 이 방법들은 상당히 많은 계산을 필요로하고, gradient explosion도 일어날 수 있다.
이러한 문제점들 사이에서, 이 논문은 denoising shortcut to fine-tune the foundational model을 만들기 위해 few-step diffusion model을 활용한 Shortcut-based Fine-Tuning (SHORTFT)방법을 도입한다. 이 기법은 denoising chain의 원래 길이를 우회하여 추론에서, 더 짧은 길이의 denosing chain을 만든다. 이 과정에서 timestep-aware LoRA를 적용하고, 이것은 추론 단계에서 비용 증가 없이 훈련가능한 파라미터의 수를 증가시켜 더 빠른 수렴과 향상된 성능을 보여준다. 또한 progressive training strategy을 통해 denoising shortcuts을 사용하여 training inference bias를 완화시킨다.
Related Work
Alignment of diffusion models
diffusion model들이 여러 조건에 더 잘 동작하도록 하기 위한 여러 가지 기법들이 존재한다.
- Fine-tuning via data augmentation
- Fine-tuning via reward models
Diffusion distillation
Diffusion model의 distillation에 대한 존재하는 방법은 주로 2가지로 분류된다.
- trajectory-preserving distillation: ordinary differential equation (ODE)에 의해 결정되는 original denoising trajectory를 보전하는 것을 목표로 하는 방법
- trajectory-reformulating distillation: 중간 trajectory steps을 무시하고 denoising의 끝 지점을 주요 supervision으로 활용하는 데 초점을 맞추는 방법
이 논문에서는 trajectory-preserving distillation 방법을 사용하고, 이 방법은 denoising shortcut에서 사용된다.
Fine-tuning few-step diffusion models
현재 연구들은 성공적으로 diffusion model들을 few step으로 fine-tuning 할 수 있도록 연구되었다. 하지만 foundational model과 대조적으로, few-step diffusion model들은 성능 저하와 distillation에 의해 capacity 감소로 부적절한 fine-tuning으로 이어질 수 있다.
SHORTFT
Shortcut-based Fine-Tuning (SHORTFT)는 trajectory-preserving few-step diffusion model에서 shortcut을 활용하여, 바로 diffusion sampling process을 통해 end-to-end backpropagation이 가능하게 한다. 이 방법은 reward function와 함께 align 하기 위한 사전학습된 diffusion model의 파라미터를 fine-tune한다.
Problem formulation
SHORTFT는 사전학습된 diffusion model들의 파라미터 θ를, 생성된 이미지에 대한 미분 가능한 reward function R(·)을 최대화하도록 fine-tuning하는 것에 초점을 맞춘다.
$Sample (θ, c, x_T)$는 prompt condition c과 함께 timestep t = T → 0에 대한 denoising process를 나타낸다. Eq.1은 **∇R(Sample(θ, c, xT), c)**을 계산하고 gradient ascent를 사용하여 푼다. 이 gradient의 계산은 denoising chain에서 multiple diffusion models을 통한 backpropagation를 필요로하고, recurrent neural networks에서 시간을 통한 backpropagation와 유사하다.
Denoising shortcut
trajectory-preserving few-step diffusion models은 본질적으로 denoising shortcut을 도입하고, denoising chain에서 high-quality와 일관된 이미지 합성을 보장하면서 융통성있는 skipping을 가능하게 한다.
Shortcut-based fine-tuning
backpropagation을 통한 naive optimization of Eq. 1는 완성된 denoision chain 구축을 필요로한다. 이 과정은 GPU VRAM안에서 각 neural layer와 각 denoising timestep과 연결된 중간 activations의 저장공간을 포함하고, 그것은 메모리 제약때문에 실행하기 힘들다. 더욱이, denosing chain은 대략 50 정도의 길이를 가져 gradient explosion 문제도 발생할 수 있다.
SHORTFT의 key insight는 denoising chain의 길이를 감소시키는 denoising shortcut 구축을 위한 trajectory-preserving few-step diffusion model의 활용이다.
Shortcut-based denoising chain
denoising shortcut을 활용하여, shortcut-based denoising chain을 구축할 수 있고 end-to-end backpropagation을 통해 the diffusion model을 fine-tune할 수 있다.
Timestep-aware LoRA
diffusion model의 모든 weight들을 다 fine-tuning시키는 것 대신에, Low-Rank Adaptation (LoRA)는 사전학습된 모델의 weight들을 보존하고 새로운 원래 모델 weight에 속한 새로운 low-rank weight matrices를 도입한다. 전체 denoising process 내내LoRA parameter들을 공유하는 것은 denoising에서 독특한 패턴이나 현상을 포착하는 걸 실패할 수 있다.
그러므로, 일반 LoRA와 대조적인, timestep-aware LoRA을 도입한다. 이 timestep-aware LoRA 설계는 추론 단계에서 computational cost 증가 없이, 효과적으로 diffusion model의 능력을 향상시키고, 훈련에서의 수렴을 가속화한다고 한다.
특히, 그림 5를 보면, 처음에 전체 denosing chain을 k개의 구간으로 분할한다. 첫 번째 구간을 제외하고, 이후 각 구간의 마지막 timestep에서 동일한 LoRA 파라미터를 공유한다. 중요한 점은, 후분 구간에서, LoRA가 없는 연속된 timestep 구간이 존재하는데, denoising shortcuts을 지원한다.
Progressive training strategy
비록 few-step diffusion model이 전체 denoising chain을 건너뛰는 denoising shortcut을 생성하는 것을 가능하게 하지만, 본질적으로 출력에서의 오류를 만들어낼 수 있다. 특히 세부적인 부분에서는 Fig. 3에 나타난 것처럼 SD의 출력과 완전히 일치할 수 없다. 이러한 불일치는 training-inference gap을 초래하여 최적이 아의 결과를 낼 수 있다. 이를 완화하기 위해 우리는 Progressive training strategy을 설계했다고 한다.
Fig. 6에 나타난 바와 같이, timestep-aware LoRA의 설계에 따라 저자들은
- SHORTFT 학습 과정을 k개의 단계로 나눈다.
- i-번째 학습 단계에서는 LoRA i부터 LoRA k까지의 가중치를 최적화한다.
- i-번째 구간 및 그 이전의 denoising process는 원래의 denosiong chain을 유지하는 반면, i-번째 구간 이후의 denoising process에서는 denoising shortcut을 도입하여 backpropagation chain의 깊이를 단축한다.
또한 runcated backpropagation기법을 사용한다. 추론 시에는, Fig. 6에 나타난 것처럼, denoising shortcut은 건너뛰고 original denoising chain을 사용하여 최종 출력 이미지를 생성한다.
Experiments
세팅이나 실험 디테일은 논문에서 확인할 수 있다.
Datasets
Human Preference Score v2 dataset (HPDv2)을 사용하였고 2 A800 GPUs로 6시간동안 훈련하였다고 한다.
Qualitative comparison
Quantitative comparison
More results
Ablation study
Conclusion
이 논문에서는, 새로운 Shortcut-based FineTuning (SHORTFT)방법을 제안한다. 이 방법은 denoising chain에서 end-to-end backpropagation을 통해 보상 함수를 사용하는 diffusion models을 aligning 시키기 위한 발전된 기법이다. 존재하는 기법들은 계산 비용과 gradient explosion 등의 문제가 있었지만, SHORTFT는 shorter denoising chains을 활용하여 fine-tuning 효율성과 유용성을 향상시킨다. 여러 다양한 실험들로 SOTA 성능을 넘는 것을 확인할 수 있다.