메타분석 시리즈: Loss re-weighting before and after the era of LLMs

들어가며

딥러닝 연구에서, 다양한 문제들이 기존 loss-function에 sample별로 특정 weight를 부과하는 방법 (Loss-reweighting)으로 풀어져 왔었다. 직관적으로 생각했을 때, 이러한 방식은 "어떤 샘플들을 중점적으로 학습시켜야 하는가?"라는 식의 물음이 이어지는 문제를 풀어내는 도구가 될 수 있었다. 따라서 저런 물음이 많이 나오게 되는 Domain-generalization 문제---어떤 데이터 도메인 (e.g. 그림, 사진, 선화, 애니메이션)을 중점적으로 학습시킬까?---, multi-task 문제---여러 task를 동시에 푸는 모델이 어떤 task (depth estimation, object-detction, segmentation)를 중점적으로 학습해야 할 까?,Imbalance-label---샘플 수가 적은 class에 대해서 중점적으로 학습해야 하지 않을까?---등을 풀어내는 도구로써 loss-reweighting이 잘 사용되어 왔다.

 

심지어, 이런식의 loss-function의 변경을 통한 학습 개선이 잘 적용되지 않는 LLM 시대 이후 (DPO 같이 아예 문제가 달라지는 걸 빼면)로도 loss-reweighting은 그 명목을 잇고 있다. 학습된 모델의 일반화 능력 (accuracy라고 생각해도 대부분의 상황에서는 무관하다)을 향상시키는 기존의 목적과는 다르게, Large-scale model (diffusion, LLM)에서는 loss-reweighting이 학습을 가속화시키는 또 다른 목적으로 사용된다. 또한, 어찌 보면 최근의 Data-curation이나 dataloader에 특정한 sampler를 붙이는 것 또한 이러한 loss-reweighting의 data-driven 한 접근이라고 볼 수도 있다 (e.g. text to image model에서, high-quality image를 중점적으로 학습시켜야 한다는 목적을 위한 curation).

 

이 메타분석 시리즈 (라고 쓰고, 캐주얼한 review-paper 라고 읽어도 무관할 듯하다)에서는, 이러한 loss-reweighting 이 사용되어 온 예시들을 리뷰하여 ["어떤 문제를 집중할지 정하기"] -> ["샘플의 중요도를 측정할 metric 정의"] -> ["해당 샘플들을 중점적으로 학습"] 이라는 형식 안에서 각자 어떠한 전략을 썼는지를 알아볼 것이다. 또한, loss-reweighting이 LLM 시대 (아마도 2022년 어딘가에서 도래한) 이후에 용도 변경되어 "학습을 가속화시키는 목적으로 사용된 케이스를 몇 가지 살펴본다. 그리고, loss-weighting이 어째서 LLM 시대 이전 (low-data regime) 에서 accuracy를 올리기 위해서 쓰였고, 이제 large-data regime에서는 학습을 가속화시키기 위해서 쓰이게 되었는지에 대한 나름의 가설을 펼쳐본다.

---

Loss-weighting before the LLM age: improve generalization

이 때의 loss-weighting 기법들은 주로 (1: Multi-task learning) 모델이 여러 task를 동시에 학습하려고 할 때 특정 task에 대해서 잘 못하거나, (2: Domain adaptation and generalization) 모델이 여러 domain의 데이터를 동시에 학습할 때---e.g. 그림, 사진, Quick-draw--- 특정 domain에 대해서 잘 못하거나 하는 문제를 해결하기 위해서 등장했다. 이 외에도 많은 분야에서 사용되었지만, 적어도 필자가 관심을 가지고 tracking 한 분야는 위 둘이다. 이러한 주제의 연구의 문제 상황들은 주로 low-data regime 혹은 long-tail distribution을 가지는 (사진은 50만장 있는데 painting은 2만 장 있음) dataset이 주어질 때 더 두드러지는데, Large-model 들에서는 모델과 데이터셋의 scaling을 통해서 그런 문제점이 어느 정도 해결되어, “성능을 높이기 위한 “ weighting 방법론은 잘 연구되지 않는 것으로 보인다.

A main concept of loss-weighting: focusing on the harder task and samples

앞서 언급했듯, Loss weighting의 기본 컨셉은, neural network 모델이 어떠한 sample 혹은 task에 집중하여 학습하여야 하나를 결정해 주는 것에 있다. 몇몇 연구에서는 집중하여야 하는 task 혹은 sample을 "모델이 배우기 어려워 하는 sample or task"로 정의하는 경우가 많다. 실제로 neural network 모델들은, 배우기 쉬운 Task에 bias 된 채로 학습을 진행한다는 연구결과가 있다---이를 Short-cut bias of neural network 혹은 spurious correlation 이라 부르기도 한다 (E.g. Geirho et al., 2020). 하지만, 특정 조정 (여기서는 loss-weighting이다)을 통해서 모델이 어려운 task에 집중하여 학습하게 만들면, 모델의 일반화 능력이 상승할 것이라는 가정에서 출발한다. 아래의 여러 배경에서의 소 사진을 Input으로 받는 image-classification이 좋은 예시가 될 수 있겠다.

어떤 classification model이 학습한 데이터가 소는 초원의 배경에, 낙타는 사막의 배경에, 북극곰은 설원의 배경에만 있었다고 해보자. 그렇다면 "사막에 있는 소"는 모델의 입장에서 매우 어려운 문제일 수 있다 (그러한 데이터가 매우 소수일것이고, 배경과 물체의 concept의 충돌이 나는 것이다.) 이러한 데이터들에 조금 더 집중해서 학습한다면, neural-network는 background를 보고 해당 물체의 class를 예상하는 것이 아니라, 실제 object를 보고 class를 추정할 것이라고 기대할 수 있다.

 

위의 소-배경 예시와 비슷한 classification 문제. Waterbird dataset에 대한 짧은 설명이다. 이미지를 결정하는 concept가 여러개가 있다고 할 때 (새의 종류, 해당 새의 배경), rare한 concept에 대해서는 모델이 예측하기 어렵다.

 

이제 우리는 모델이 배우기 어려운것에 집중하면 어떠한 장점이 있을 수 있는지에 대해서 알아보았다. 그렇다면, 실제로 (1) 모델이 어떤것을 어려워하는가 (2) 그것을 어떻게 정량화하는가 에 대한 문제가 남아있다. 이때, 어떤 요소가 모델이 배우기 쉬운지 어려운지를 결정하는지에 대해서는 여러 연구에서 정의하는 바가 다르다. 여기서는 모델이 어떤 요소들을 "배우기 쉽다" 여겨서, 해당 feature에 bias가 생기는지를 세 가지로 분류해 본다.

"ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness", CNN은 texture를 더 중점적으로 배우고 이들에 bias를 가진채로 학습한다는 연구.

• (예시 1: frequency bias) Vision → CNN 모델이 high-frequency input 등에 bias 되어 학습하기 때문에, texture에 따라 image를 분류하는 현상이 있다
• (예시 2: task bias - multi-task learning) 자율주행을 위한 모델이 segmentation, depth estimation을 동시에 학습한다고 할 때, 특정 task에 너무 과도하게 집중하여 다른 task를 잘 못하는 경우가 있다. (Label이 여러 도메인을 가짐)
• (예시 3: Domain bias - domain generalization) 이미지를 생성하거나 분류하는 모델이 있을 때, 사람 얼굴 image에서는 제대로 동작하지 않지만 풍경 등에서는 잘 동작한다. (Input data가 여러 도메인을 가짐)

Domain-generalization 문제의 예시. 어떤 classification model에서 painting, sketch, cartoon 형태의 도메인을 가지는 이미지는 많으나, 실제 사진에 대한것은 매우 적거나 없다고 가정하면, 실제 사진 input에 대해서는 classification model이 풀기 어려울 수 있다.

 

이러한 bias는 해당 task에 대한 데이터가 작아서 발생할 수도 있고 (위의 waterbird example에서와 같이), 모델의 특성 자체에서 기인할 수도 있다 (CNN 구조 자체에 high-frequency bias 가 있다던지). 혹은, 학습을 더 진행하게 되면 어려운 task에 대해서 잘 수행할 수도 있다.

 

모델의 Bias를 해결하는 방법론은, 모델이 풀기 어려운 task에 집중해서 학습하게 하는 계열들이 많다. 이 중 가장 직관적인 것이 loss-weighting 전략이다. 즉, “Hard task “인 sample에 대해서 높은 loss-weight을 걸어주면서 학습하는 것이다. 그렇다면 어떤 것들이 모델 입장에서 “어려운” task일까? 어떻게 그 어려움을 정량화하고, 형식화하여 이를 활용해 loss-function을 설계할 수 있을까? 이 질문이 해당 문제를 풀려고 하는 여러 연구들이 알아내고자 하는 것이었으며, 이 질문에 잘 답하는 것이 training efficiency와 generalization capability를 향상하는 열쇠가 되었다. 

 

다음 장들에서는 loss-weighting이 자주 등장하는 두 연구주제 multi-task learning과 domain-generalization에 대해서 “어려운 task”를 어떻게 찾아서 weighting을 하는지에 대한 예시를 몇 가지만 간단히 소개한다. multi-task learning의 경우, 앞선 설명처럼 "모델이 풀려고 하는 여러 task 중 어떤 task가 어려운 task 일까?"에 대한 물음에 답하려 하고, domain-generalization의 경우 "모델이 예측하려는 여러 domain의 data 중 어떤 domain이 어려운 domain일까?"에 대한 물음에 답하려 한다. 아래에서, 두 연구 주제에서 대표성을 띄는 세 가지 논문들에 대해서 어떻게 어려운 task를 정의했는지, 그리고 그 어려운 task에 집중하도록 하였는지를 정리한다.

 

# (NOTE; Informal): LLM 이전 파트에 해당하는 아래 예시들은 매우 간단하게만 소개한다. 이 파트는 마치 논문에서의 related works에 가깝다. 즉, 이 글의 구성중 아마 제일 졸리는 파트일 것이다. 이 이후에 본격적인 이야기를 시작할 테니 좀 참아달라.

Loss-weighting and multi-task learning

Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics

• (어려운 task에 대한 정의): 예측 불확실성 (Uncertainty)이 높은 task → 모델이, 나 이거 확신이 없는데?라고 말하는 task
  • (Classificaiton의 경우) Softmax의 entropy가 높은 것
  • (Regression의 경우) Output을 Gaussian distribution의 파라미터로 바꾸었을 때, entropy가 높은 것.
• 방법론: 높은 uncertainty를 가지는 sample의 loss-weight을 높게 설정함, 이는 매 iteration마다 바뀌게 됨

GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks

Multi-task learning에서 한 때, 많은 연구들의 baseline으로 사용된 기법이다. Gradient의 imbalance가 task bias를 만들어 주는 것이라 가정해서, gradient 자체를 normalize 해 준다. 이때, 실질적으로 loss를 reweighting 해주는 것으로 gradient를 normalize 시킨다.)

GradNorm은 특정 task마다의 gradient의 imbalance로 부터 특정 bias가 생긴다고 생각하여, 이들을 그림의 오른쪽처럼 적절히 normalize 시켜주면 문제를 해결할 수 있을것이라 가정한다.

• (어려운 task에 대한 정의): 특정 task의 loss로부터 생성되는 gradient의 L2 norm이 지속적으로 낮은 task.
• 방법론: 위 정의에 따라, gradient를 여러 task별로 normalize 시킴

Loss-weighting and domain generalization

Domain generalization: 데이터의 source 혹은 domain이 여러 개인 ({그림, 사진} 혹은 {영어, 한국어, 일어}) 문제이다. (https://bayesian-bacteria.tistory.com/5)

Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization

Distributionally Robust Optimization (DRO) 기법은, 일반적인 학습 방법론인 ERM (그림의 왼쪽, Empirical Risk Minimization: 내가 가지고있는 empricial distribution 상에서 각 데이터 샘플이 1/n 만큼의 "중요도"를 가짐) 에서 벗어나서, 각 group (domain과 비슷하게 생각해도 좋다.) 별 loss에서 가장 못하고 있는 group에 대해서 계속 최적화를 시켜주는것을 기본적인 아이디어로 한다 (그림의 오른쪽).

• (어려운 task에 대한 정의) Dataset이 여러 distribution에서 왔다고 가정했을 때, expected loss가 가장 높은 distribution 에 대한 task
  • Informal 한 예시: 한국어, 영어, 일어의 모든 data를 가지고 있다고 가정했을 때, average loss가 가장 높은 언어에 대한 task
• 방법론: 여러 domain 중에서 가장 loss가 높은 domain을 매 iteration마다 최적화

 

위에서 언급한 세 가지의 논문들은 각자의 어려운 task에 대한 정의를 통해서 loss-reweighting으로 풀고자 하는 문제 (특정 task에 대한 bias, 특정 domain에 대한 bias)를 성공적으로 해결해서 모델의 generalization capability를 높였었다. 이제 아래에서는, LLM과 Diffusion 등 좀 더 modern 한 모델 구조와 task에서 이러한 loss-reweighting이 어떻게 다른 목적인 training 가속화를 이루는지, 그리고 왜 LLM 시대 이전과 이후의 연구가 달성하고자 하는 목적이 이렇게 달라졌는지를 다뤄본다.

---

Loss-weighting before the LLM age: boost the training

Why does loss-weighting work? And why does it work differently in the past compared to now?

앞서 언급하였듯, loss-reweighting은 LLM 이전 시대 즉, 데이터셋과 모델의 규모가 작을 때는 generlization capability를 증가시키는 방식으로 활용되었고, 근래에 들어서는 학습의 수렴속도를 빠르게 하는데에 사용되었다. 여기서는 같은 방법론 (loss-weighting) 이 왜 다른 목적으로 사용되어 왔을까? 에 대해서 논하려고 한다. 

 

(주의) 여기서의 썰은 어떤 reference가 있는 것이 아니고 순전히 필자의 뇌피셜 및 사고 실험들이다.

 

이를 위해서, 위 그림과 같이 특정 모델 parameter $\theta_0$ 이 loss-landscape에서 3가지 task (혹은 domain) easy-task, hard-task, super-hard-task 에 대한 구덩이 (basin) 쪽으로 최적화가 되어 나가는 상황을 가정해 보겠다. 여기서, 모델의 사이즈와 데이터가 큰지 작은지, loss-weigthing을 하는지 안 하는지에 따라서 optimiziation trajectory가 어떤 방식으로 생성되는지에 대한 사고실험을 해볼 것이다. 먼저 아래와 같은 가정을 깔고 위 그림에서의 케이스별로 어떻게 학습이 진행될 것인지를 생각해 보도록 하자.

Assumptions

• (Assumption 1) Neural network는 쉬운 것부터 학습하려고 한다. (앞서 short-cut bias, spurious-correlation 등으로 논했던 연구주제에서 언급한 바와 같이.)
• (Assumption 2) 여러 task에 대한 mode가 loss landscape에 있다고 가정해 보자. 이때, 만약 task 간의 괴리가 크면 (task 간 mode가 멀리 떨어져 있다던지.), 어떤 task에 최적화된 mode에 들어간 이후로는, 다른 mode에 들어가기 어렵다.

위 가정과 함께, 아래와 같은 세 가지 케이스에서의 model-weight optimization 시나리오를 생각해 보면, loss-weighting은 적은 데이터셋에서는 "도달 불가능한 sweet-spot에 도달할 수 있게 함"으로서 일반화 능력을 향상하고, 큰 데이터셋과 큰 모델에서는 "도달은 가능하지만, 최단경로로 도달하지는 못하는 어떤 sweet spot을 좀 더 빨리 도달할 수 있게 함"으로 학습을 빠르게 해 줄 수 있을 것이라 생각할 수 있다. 아래에서 좀 더 상세히 논해보자.

Case 1. 적은 데이터셋과 적은 모델.

이 상황에서는, 데이터셋이 적고 모델이 작기 때문에 (Double-descent regime에 따라서, model size가 크면 클수록 일반화가 쉽다고 가정) 개별적인 3개의 task가 서로 꽤나 떨어져 있다. 처음의 모델은 easy-task에 bias 된 채로 최적화를 진행해 나가다가, 점차. hard-task를 어렵게 학습해 나가지만, 끝끝내 super-hard-task로 가기 위한 barrier는 통과하지 못했다. 

(Assumption on the effect of loss-weighting with small-dataset) 만약, loss-weighting을 통해서 애초부터 super-hard-task 방향으로 학습이 진행되었다면, 도달 불가능한 sweet-spot에 도달할 수 있었을 것이다. 

 

Case 2. 큰 데이터셋과 큰 모델.

모델과 데이터셋이 커짐에 따라, 서로 다른 세 task는 어느 정도 인접하게 된다. 따라서, 모델이 easy-task 에 bias 된 채로 시작하더라도, 끝끝내 hard-task 와 super-hard-task 에 도달할 수 있게 된다 (basin 사이의 장벽이 거의 없다.) 하지만, 모델이 easy-task 에 bias 된 채로 학습했었기 때문에, 최단경로로 sweet-spot (모든 task를 잘하게 되는 공간)에 도달하지는 않았다.

 

Case 3. 큰 데이터셋과 큰 모델, 그리고 loss-reweighting

loss-weighting을 통해, 처음부터 어느 정도 hard-task와 super-hard-task 를 잘하는 방향으로 학습하게 된다. 따라서, easy-task에 대한 optima를 거치고 sweet-spot에 도달해서 (즉, 우회로를 통해 optimized 되어서) 시간이 오래 걸렸던 case 2에 비해서, 더 최단거리로 sweet-spot에 도달할 수 있게 된다. 

(Assumption on the effect of loss-weighting with small-dataset) loss-weighting은 도달은 가능하지만, 최단 경로로 가지는 못하는 어떤 sweet-spot에 좀 더 빨리 도달하게 한다.

 

완전 제멋대로 입맛대로 가설이지만, 위와 같은 상황을 생각하면 loss-weighting이 왜 과거에는 일반화 능력을 향상하는데 중점을 두었으며, 최근에는 해당 objective에서 벗어나서 학습을 더 빨리 하는것에 초점을 맞춰서 연구되는지에 대해서 어느정도 납득해볼 수 있다. 이제 다음 장에서는, 실제로 loss-weighting 으로 학습을 가속화 시키는데 사용된 두 연구 사례들을 살펴볼 것이다.

 

Case study (LLM): DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

• 어려운 task에 대한 정의: 어떤 small-model에서 loss가 높은 domain (task) 은, 어려운 domain으로 생각한다.

DoReMi는 아주 직관적으로 어려운 task를 정의하려고 했다. 바로 loss가 높게 나오는 특정 도메인에 대해 token을 예측하는 것을 어려운 task라고 보는 것이다. 예를 들어, 2B 짜리 language-model을 가져왔을 때, 유독 Coding task에 대해서 높은 loss를 보인다면, coding task를 더 중점적으로 학습하고 싶은 것이다. 구체적으로는 위 그림처럼 3-steps 전략을 사용한다:

1. Small model을 일반적인 방법으로 학습한다 (Reference model)
2. 또 다른 small model (proxy-model)을 학습하는데, 특정 domain에 loss가 높게 나오면, 다음 iteration부터는 해당 domain에 더 집중해서 (weight을 높여서) 학습한다.
3. proxy model을 학습할 때 사용했었던 “Domain weight “로, 최종적으로 학습할 large-model을 학습한다.

 

어려운 domain에 더 높은 loss-weight을 주고, 해당 weight을 large-model 학습에서 사용하겠다는 전략이다. 작은 모델에서의 domain-weight이, 해당 domain의 난이도를 잘 반영해 주겠다는 믿음이 보인다. 그렇다면, 이러한 domain-weight은 어떻게 학습되는지를 결정해야 한다. (재미있게도 이런 domain-weight을 결정하는 것은 앞서 말했던 LLM 시대 이전의 DRO 논문에서 제시된 것과 닮아있다: Group-DRO Section 5)

Step 1에서 학습한 reference model을 활용하여, Step 2에서 학습하는 proxy model의 domain weight $\alpha_t$를 업데이트 하는 방식.

 

위 그림에서처럼, 시간 $t$ 에서 Domain weight $\alpha_t \in \mathbb{R}^{k}$ 은 domain weight의 logit 에 해당하는 $\lambda_t[i]$ 에 의해 생성된다. 이 때, domain-weight-logit $\lambda_t$ 는 특정 도메인 $D_i$에 대해서, 현재 학습중인 proxy model의 loss $l_{\theta_{t-1}, j}(x)$ 와 step 1에서 학습했던 reference model의 loss $l_{\text{ref},j}(x)$ 의 차이 (붉은 박스)가 클수록 큰 값을 가진다. 근데 왜 하필, 그냥 loss가 높은 sample이 아니고 reference model의 loss와의 차를 이용한 걸까? 그냥 Loss $l_{\theta_{t-1},j}(x)$ 가 높은 domain의 weight을 증가시키면 안 되는 걸까?

 

 

이것은, 본질적으로 어려운 샘플들을 제외하고 싶어서이다. 즉, reference model loss가 매우 높은 샘플들은 "애초에 배우기 어려운" 도메인으로 생각하는 것이다. 대표적인 예시가 noisy-label이다. noisy label은, loss는 매우 높게 나와 배우기 어렵지만, informative 한 샘플은 전혀 아니다. (중간 광고로 이러한 noisy-label과 informative한 sample 간의 충돌을 다룬 필자의 arxiv 논문---이라 쓰고 ICLR 광탈이라고 읽겠다---을 링크건다.) 이러한 uninformative-but-hard-to-learn 한 샘플들은 reference model의 Loss도 높을 것이기 때문에, 위 알고리즘 상에서는 적절히 걸러진다.

 

실제로 우측 상단그래프에서 "Hardest"가 단순히 high-loss sample에 대해서 domain-weight을 높인 방법인데, uninformative-but-hard-to-learn 한 샘플을 학습한 DoReMi 에 비해 underperform 하는 것을 볼 수 있다.

 

이렇게 얻어낸 domain weight $\alpha$ 로 실제 우리가 학습하고 싶은 large-model을 학습하면 DoReMi의 3-step 학습이 종료된다. 그렇다면, 이것은 얼마나 효과적이었을까? 얼마나 효과적으로 모델 학습을 가속화시켰을까? 오른쪽 그림에서 처럼, 동일 수준의 Exact token match를 도달하는 데에 최대 2.6배의 가속이 있었다 (필요 step 수에 있어서). 그리고, 논문에서 주요 contribution으로 생각하진 않으나 마치 LLM 이전 시대의 논문들처럼 accuracy 측면에서도 향상이 되었음을 보인다. 아마도 계속 학습을 시키면 결국 같은 성능에 도달하기 때문일 것이라 생각된다.

 

Case-study (Diffusion): Efficient Diffusion Training via Min-SNR Weighting Strategy

• 어려운 task에 대한 정의: Diffusion의 early step (noise가 큰 샘플)을 denoising 하는 것이 어렵다.

Diffusion에서는 여기서 소개하는 Min-SNR gamma처럼 "Diffusion의 timestep 별로 어려운 task가 있다"라고 생각한다. 실제로 논문에서는 multi-task learning 문제로 diffusion을 formulation 할 수 있다고 abstract에서도 언급한다. 

보통 Diffusion model에서는 위와 같은 특정 timestep t의 input $x_{t}$를 모델 ($f_\theta$) 의 input으로 받아 노이즈가 한 단계 제거된 $x_{t-1}$ 을 output으로 내뱉어준다: $f_{\theta}(x_t, t) \approx x_{t-1}$. 이에 대한 loss는 단순히 실제 $x_{t-1}$ 과의 Mean-square-error (MSE) loss를 설정하는 것이 보통이다: $\|f_{\theta}(x_t, t) - x_{t-1}\|^2_2$. 이때, 생각해 보면 별로 noise가 끼어있지 않은 이미지 (e.g. $x_3$) 에서 살짝 denoise를 하는 것은, 매우 심한 noise에서 denoise 하는 것보다 꽤 쉬운 task이다. 이미 사람 얼굴의 형상을 가진 이미지에서 조금 더 texture 만 보강해 주는 것과, 아예 처음부터 어떻게 생겼는지도 모르는 얼굴의 형상을 만들어내야 하는 것을 생각해 볼 수 있다.

 

앞에서 언급한 것처럼 논문의 noise-level이 높은 이미지를 중점적으로 학습해야 한다는 주장의 rationale를 정리하면 이러하다:

 

• (Rationale 1; informal) 이미지를 denoising 한다고 할 때, 원래 이미지가 거의 훼손된 pure-noise에 가까운 이미지가 더 denoising 하기 어렵다. (오른쪽 그림에서 처럼, t가 클 때 error가 더 크다)
• (Rationale 2“) 기존의 multi-task learning에서의 관측처럼, 특정 task(timestep)에 diffusion을 finetune 시키면 해당 task에 대해 더 denoising을 잘하게 되는 연구가 있다. 다만, 다른 task에서는 잘 못하게 된다. 따라서, 적절한 weighting을 통해 이 Multi-task learning 문제를 풀면, 모든 task에 대해 잘 하게 만들 수 있을 것이다. 

따라서, Min-SNR Gamma는 "noise level이 높은 상황을 중점적으로 학습" 하도록 한다. 더 정확히 말하면, noise-level이 낮은 것에 과도하게 집중하지 않도록 한다. 사실 이러한 원칙은 Min-SNR gamma 말고도 여러 연구들에서 적용하였는데, 모두 일반적인 학습보다 더 높은 성과를 거뒀다. 아래에서의 왼쪽 그래프의 "Min-SNR", "P2", "Ours (Debiasing the training diffusion models)" 가 해당 연구들이며, 그래프의 y축이 timestep $t$ (x축)에 따라 걸리는 loss-weight이다.

이런 단순한 방식이 충분한 training-efficiency boost를 해 주었을까? 우선 논문에서 제시된 바로는 그러하다, 수치적으로는 아래와 같다.

• Efficient diffusion training via min-str weighting strategy에서는 vanilla diffusion에서 600K iteration에 도달하는 FID를 400K 만에 도달하였다 (ImageNet 256x256)
• 앞서 언급한 유사한 연구인 Debias the training of diffusion models에서는 1000K iteration의 FID를 400K 만에 도달하였다. (FFHQ dataset - human faces)

---

Discussion

이번 메타 분석에서는 loss-weighting이 LLM 시대 이전에는 generalization 능력을 강화시키고, LLM 이후에는 training efficiency (적은 step으로 높은 성능 도달)를 높이는 방식으로 사용됨을 보았다. 그리고, 같은 방법론이 왜 large-model에서는 efficiency를 증가하는데 도움이 되는 것이 main-contribution이고, small-dataset small-model에서는 generalization 능력을 향상하는 것이 main-contribution이었는지에 대한 추측을 해보았다.

 

아래에서, loss-weighting의 과거와 현재의 응용 사례로부터 나오는 필자의 몇 가지 질문을 던지면서 본 글을 마무리하려고 한다.

DoReMi의 사례로부터: 해당 연구들은 오히려 Data-centric AI의 중요성을 강조하는가?

• 처음부터 Domain 혹은 Task 간의 balance를 가지고 있는 양질의 데이터셋이나, noisy 하거나 ambiguous 한 sample을 제거하면 DoReMi의 사례와 동일한, 혹은 더 뛰어난 성능을 보일 것으로 기대해 볼 수 있다.
• 다만, 위처럼 명시적인 data-selection 자체가 어려운 경우 본 페이지에서 제시되는 방법론들을 시도해 볼 수 있을 것.

Large-model regime에서는 weighting strategy가 좋을까?

• Diffusion: 지금까지 잘 사용되는 연구사례들은, (1) loss-weighting이 step별로 바뀌거나 하지 않고; (2) loss-weighting에 모델이 관여한다든지 (DoReMi의 사례처럼) 하지 않는다. 이런 방식을 Diffusion에는 적용 불가능한가? 또한, 요즘 자주 사용되는 flow-matching에서는 적절한 weighting이 있을까? (SD3에서 이미 보여주었지만, 더 향상된 무언가가 있을까?)
• Large-models: High quality, informative, and hard to train 한 샘플들을 찾아내는 더 좋은 방식은 없는가?

Weighting strategy는 과연 Large-model의 성능 자체도 올려줄 것인가?

• 특정 Accuracy에 빨리 도달한다를 넘어서, 과거의 loss-weighting 응용 사례처럼, 무한한 시간을 학습했다고 할 때 실제 “accuracy” 도 올라갈 것인가? 그렇게 학습된 모델에서는 최종 accuracy에 domain 간 trade-off 관계가 있을까? (특정 task를 잘 해내지 못한다던지..) 혹은 pareto-front를 뚫어낸 어떤 모델이 학습될까?