[Github]
[arXiv](2023/10/22 version v1)
Abstract
Distillation 없이 consistency model을 직접 학습하는 consistency training의 개선된 기법 소개
Consistency Models
잔차 연결을 통해 CM을 피라미터화. σ = noise level
CM loss:
N이 작을수록 Δt가 크다.
fθ = student network, fθ̄ = teacher network라고 한다. Teacher network는 EMA로 업데이트된다.
일관성 증류(CD)의 경우 확산 모델 s를 통해
Consistency Training (CT)의 경우 데이터 x와 가우시안 노이즈 z를 활용하여
훈련 후 가우시안 노이즈 z에서 시작하여 한번에 생성하거나
다단계로 생성할 수 있다.
Improved Techniques for Consistency Training
수정 사항 한눈에 보기:
Weighting functions, noise embeddings, and dropout
1. 낮은 noise level에 더 높은 가중치를 할당한다.
이는 정보가 더 많이 보존되어 있는 낮은 noise level에서의 오류가 높은 noise level의 손실에 영향을 미칠 수 있다는 이유에서다.
2. 기존 CM에서 noise level 임베딩에 사용되었던 fourier embedding layer의 민감도를 positional embedding과 비슷한 수준으로 낮추어 훈련 안정성을 향상한다. 근데 CM 논문에서 fourier embedding에 대한 내용은 없었는데
3. CM의 과제가 DM 보다 더 어려우므로 더 어려운 과제에서는 과적합이 적을 거라는 직관에서 기존에는 dropout을 사용하지 않았지만 사용했을 때 샘플 품질이 더 좋았으므로 dropout을 사용한다.
Removing EMA for the teacher network
Teacher network에서 µ = 0으로 하여 EMA를 제거하면 품질이 크게 향상된다. 이론적 증명은 논문에
Pseudo-Huber metric functions
하이퍼 피라미터 c를 조절하여 outlier에 대한 민감도를 조절할 수 있어 L2 loss보다 강건한 pseudo-huber loss를 사용한다.
Improved curriculum for total discretization
Discretization step N이 클수록 샘플 품질이 향상되며 편향이 줄어들고 분산이 증가한다. 편향과 분산의 균형을 고려하여 N의 상한선을 1281로 설정했다.
또한 일정한 training iteration k 마다 N이 2배로 증가하는 exp curriculum을 사용한다. (노란색)
Improved noise schedules
기존의 noise schedule은 높은 noise level이 더 높은 확률을 가져 더 많은 data point가 생성된다. 이는 낮은 noise level의 오류가 후속 noise level의 오류 누적을 유발한다는 직관에 반하며, 다음과 같이 이산화 된 로그 정규분포에 따라 noise level을 샘플링한다.
