Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization

이 논문은 SAM 의 단순성 편향 감소 특성을 활용하여 학습 초기에 데이터 클러스터링과 업샘플링을 수행하는 'USEFUL' 방법을 제안함으로써 다양한 아키텍처와 데이터셋에서 기존 최첨단 성능을 달성하는 것을 증명합니다.

핵심

이 논문은 **"인공지능 (AI) 이 더 똑똑하게 학습하도록, 학습 데이터의 구성을 살짝 바꿔보자"**는 매우 흥미로운 아이디어를 제시합니다.

기존의 AI 연구는 "더 좋은 알고리즘을 만들자"거나 "더 많은 데이터를 모으자"는 데 집중했는데, 이 논문은 **"이미 있는 깨끗한 데이터를 어떻게 섞어서 가르치느냐"**가 핵심이라고 말합니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

🎓 비유: "수학 선생님과 편식하는 학생"

상상해 보세요. AI 모델은 수학 공부를 하려는 학생이고, 학습 데이터는 수학 문제집입니다.

1. 문제점: "쉬운 문제만 먼저 푸는 습관" (단순성 편향, Simplicity Bias)

대부분의 학생 (기존 AI 학습 방법인 SGD) 은 문제를 풀 때 가장 쉬운 문제부터 먼저 풉니다.

• 쉬운 문제 (Fast-learnable features): "고양이 사진이니까 고양이"라고 바로 알 수 있는, 귀가 뚜렷하고 배경이 깨끗한 사진.
• 어려운 문제 (Slow-learnable features): "고양이"지만 귀가 가려져 있거나, 어두운 그림자 속에 있거나, 배경이 복잡한 사진.

학생은 초반에 쉬운 문제만 계속 풀면서 "아, 고양이 = 귀가 뾰족한 것"이라고만 기억합니다. 나중에 어려운 문제를 마주치면, "귀가 가려졌으니 이건 고양이가 아니야"라고 잘못 판단하게 됩니다. 이를 논문에서는 **'단순성 편향 (Simplicity Bias)'**이라고 부릅니다.

2. 해결책: "SAM 이라는 새로운 학습법"

최근 등장한 **SAM(SAM)**이라는 학습 방법은 이 학생에게 "쉬운 문제만 풀지 말고, 어려운 문제도 골고루 보라"고 가르칩니다.

• SAM 은 문제를 풀 때 조금 더 넓은 시야를 가지고 접근합니다.
• 그 결과, 쉬운 문제와 어려운 문제를 동시에, 균형 있게 학습하게 되어 시험 점수 (일반화 성능) 가 훨씬 좋아집니다.
• 하지만 SAM 은 학습 시간이 두 배로 걸린다는 단점이 있습니다.

3. 이 논문의 혁신: "USEFUL" (데이터 섞기 전략)

이 논문은 **"SAM 처럼 똑똑하게 학습하게 하려면, 데이터 자체를 살짝 바꿔보자"**고 제안합니다.

**"USEFUL"**이라는 방법의 핵심은 다음과 같습니다:

1. 초반 학습 관찰: 학생이 문제를 풀기 시작할 때, 어떤 문제들이 가장 먼저 쉽게 풀리는지 (쉬운 문제) 확인합니다.
2. 그룹 나누기: "아, 이 문제들은 너무 쉬워서 이미 다 풀었구나"라고 판단되는 문제들을 그룹 A로, "아직 헷갈려 하는 어려운 문제들"을 그룹 B로 나눕니다.
3. 한 번만 복사해서 더 많이 보기: 그룹 B(어려운 문제) 의 문제들을 2 배로 복사해서 문제집에 다시 넣습니다. (그룹 A 는 그대로 둡니다.)
4. 다시 시작: 이렇게 섞인 새로운 문제집으로 다시 공부를 시작합니다.

결과:
학생은 이제 어려운 문제 (그룹 B) 를 훨씬 더 자주 마주치게 됩니다. 덕분에 "고양이 = 귀가 뾰족한 것"이라는 단순한 규칙만 외우지 않고, "어두운 그림자 속 고양이" 같은 복잡한 특징도 함께 배우게 됩니다.

---

💡 왜 이것이 중요한가요?

• 기존 방식의 한계: 단순히 데이터를 더 많이 모으거나, 모델을 더 크게 만드는 것만으로는 한계가 있었습니다.
• 이 방법의 장점:
  • 비용 절감: SAM 을 쓰지 않아도 (기존 SGD 를 쓰더라도) 비슷한 효과를 낼 수 있습니다.
  • 단순함: 복잡한 알고리즘을 새로 짜는 게 아니라, 데이터를 한 번만 섞어서 다시 학습시키는 아주 간단한 전략입니다.
  • 성능 향상: CIFAR-10, ImageNet 같은 유명한 이미지 인식 테스트에서 기존 최고 기록 (State-of-the-Art) 을 갱신하는 성과를 냈습니다.

🚀 요약하자면

이 논문은 **"AI 가 편식하지 않고, 모든 유형의 데이터를 골고루 배우게 하려면, 학습 초기에 '어려운 데이터'를 조금 더 많이 보여주고 학습을 다시 시작하면 된다"**는 것을 증명했습니다.

마치 요리사가 처음부터 모든 재료를 다 섞지 않고, 가장 익히기 어려운 채소부터 먼저 많이 넣고 요리를 시작하면, 결국 훨씬 더 맛있는 요리를 완성할 수 있는 것과 같은 원리입니다.

이 방법은 USEFUL (UpSample Early For Uniform Learning) 이라는 이름으로, AI 가 더 똑똑하고 편견 없이 세상을 이해하는 데 도움을 줄 것입니다.

기술 요약

1. 문제 제기 (Problem)

• 배경: 과매개변수화 (overparameterized) 된 신경망은 훈련 데이터에 대한 영구 오차 (zero training error) 해가 여러 개 존재하며, 이 중 일반화 성능이 우수한 해를 찾는 것이 핵심 과제입니다.
• 단순성 편향 (Simplicity Bias): 경사 하강법 (GD, SGD) 은 훈련 초기에 '쉬운' 특징 (fast-learnable features) 을 먼저 학습하고, 나중에 '어려운' 특징 (slow-learnable features) 을 학습하는 경향이 있습니다. 이는 모델이 데이터의 복잡한 패턴을 충분히 학습하지 못하게 하거나, 훈련 후반부에 과적합 (overfitting) 을 일으켜 분포 내 (In-Distribution, ID) 일반화 성능을 저하시킬 수 있습니다.
• 기존 접근법의 한계: 기존 연구들은 주로 더 나은 최적화 알고리즘 (예: SAM) 이나 데이터 필터링에 초점을 맞췄습니다. 하지만 깨끗한 훈련 데이터의 분포를 변경하여 최적화 과정이 더 균일한 특징 학습을 하도록 유도할 수 있는지에 대한 질문은 아직 열려 있었습니다.
• 목표: 훈련 데이터의 분포를 변형하여, 최적화 알고리즘이 SAM(SHARPNESS-Aware Minimization) 과 유사하게 다양한 특징을 더 균일한 속도로 학습하도록 유도함으로써, 분포 내 일반화 성능을 극대화하는 방법을 제안합니다.

2. 방법론 (Methodology)

A. 이론적 분석: SAM vs GD

• 이론적 설정: 2 층 비선형 CNN 과 특정 데이터 분포 (빠르게 학습 가능한 특징과 느리게 학습 가능한 특징이 섞인 형태) 를 가정합니다.
• 주요 발견:
  • GD: 훈련 초기에 빠르게 학습 가능한 특징만 학습하고, 느린 특징은 거의 무시합니다.
  • SAM: GD 에 비해 느린 특징과 빠른 특징을 더 균일한 속도로 학습합니다. 이는 SAM 이 '단순성 편향'에 덜 취약함을 의미하며, 결과적으로 더 평탄한 (flatter) 최소점을 찾아 일반화 성능이 우수해집니다.
• 핵심 통찰: 훈련 초기에 모델 출력 (model output) 을 분석하면, 빠르게 학습된 특징을 가진 샘플과 그렇지 않은 샘플을 클러스터링하여 구분할 수 있습니다.

B. 제안 알고리즘: USEFUL (UpSample Early For Uniform Learning)

이론적 통찰을 바탕으로 제안된 방법은 다음과 같은 3 단계로 구성됩니다.

1. 분리 (Separation): 훈련을 초기 (예: 5~10% 에포크) 에 진행한 후, 각 클래스 내의 샘플을 모델의 출력 (last-layer activation) 을 기준으로 k-means 클러스터링을 수행합니다.
   • Cluster 1 (Fast-learnable): 평균 손실이 낮고 모델이 쉽게 예측하는 샘플들.
   • Cluster 2 (Slow-learnable): 평균 손실이 높고 모델이 학습하기 어려운 샘플들.
2. 업샘플링 (Upsampling): '느리게 학습되는 특징'을 가진 Cluster 2 의 샘플들을 한 번만 (One-shot) 업샘플링합니다 (예: 2 배 증가). 이는 훈련 데이터 분포를 변경하여 느린 특징의 학습 속도를 가속화합니다.
3. 재훈련 (Retraining): 수정된 데이터 분포에서 모델을 처음부터 다시 훈련합니다.

• 특징: 동적 가중치 조정 (dynamic reweighting) 이 아닌, 단 한 번의 분포 변경으로 고정된 분포에서 훈련을 수행하여 계산 효율성을 높였습니다.

3. 주요 기여 (Key Contributions)

1. 이론적 증명: 2 층 CNN 에 대한 엄밀한 분석을 통해 SAM 이 GD 보다 특징을 더 균일하게 학습함을 수학적으로 증명했습니다.
2. 새로운 데이터 변형 기법 (USEFUL): 단순성 편향을 완화하기 위해 훈련 초기에 데이터 분포를 변경하는 새로운 방법을 제안했습니다. 이는 SAM 의 장점을 SAM 을 사용하지 않는 SGD 기반 모델에도 적용 가능하게 합니다.
3. State-of-the-Art (SOTA) 성능 달성:
   • CIFAR-10, CIFAR-100, STL-10, CINIC-10, Tiny-ImageNet 등 다양한 데이터셋에서 ResNet, VGG, DenseNet, ViT 등 다양한 아키텍처에 적용하여 SOTA 성능을 기록했습니다.
   • 특히, SAM + TrivialAugment + USEFUL 조합은 CIFAR-10 에서 ResNet18 기준 SOTA 정확도를 달성했습니다.
4. 범용성: SGD, SAM, ASAM 등 다양한 최적화 알고리즘과 TrivialAugment 와 같은 데이터 증강 기법과 호환되며, 노이즈가 있는 레이블이나 장꼬리 (long-tail) 분포 데이터에서도 성능 향상을 보였습니다.

4. 실험 결과 (Results)

• 일반화 성능: USEFUL 을 적용한 SGD 는 기존 SAM 보다도 낮은 테스트 오차를 기록했습니다. (예: CIFAR-10 에서 SGD+USEFUL 은 SAM 보다 성능이 우수함).
• 학습 동역학 분석:
  • Forgetting Score: USEFUL 은 훈련 중 샘플이 올바르게 예측된 후 다시 틀리는 경우 (forgetting) 를 줄여, 모델이 느린 특징을 더 일찍 그리고 안정적으로 학습함을 보여줍니다.
  • Hessian Spectrum: USEFUL 은 SGD 보다 더 평탄한 최소점 (flatter minima) 을 찾으며, Hessian 의 최대 고유값과 벌크 (bulk) 를 감소시킵니다.
  • L1 Norm: SAM 과 유사하게 더 희소한 (sparser) 해를 찾습니다.
• 비용 효율성: SAM 은 훈련 시간을 2 배 증가시키지만, USEFUL 은 데이터 분포를 한 번 변경하고 재훈련하므로 SAM 보다 훨씬 적은 계산 비용 (약 1.3~1.6 배) 으로 유사하거나 더 나은 성능을 달성합니다.

5. 의의 및 결론 (Significance)

• 패러다임의 전환: 기존에는 "더 좋은 최적화 알고리즘"이나 "더 좋은 모델 구조"를 찾는 데 집중했다면, 이 논문은 **"훈련 데이터 분포 자체를 지능적으로 변형"**하여 최적화 과정의 편향을 교정할 수 있음을 증명했습니다.
• 단순성 편향의 해결: 단순성 편향이 분포 내 (ID) 일반화에도 해롭다는 점을 최초로 이론적으로 증명하고, 이를 해결하는 실용적인 방법을 제시했습니다.
• 실용성: 복잡한 최적화 알고리즘 변경 없이, 기존 파이프라인에 쉽게 통합할 수 있으며 (한 번의 클러스터링과 업샘플링), 다양한 아키텍처와 최적화 기법과 결합하여 즉각적인 성능 향상을 제공합니다.

요약하자면, 이 논문은 SAM 의 우수한 일반화 성능이 '단순성 편향'을 줄이고 특징을 균일하게 학습하기 때문이라는 통찰을 바탕으로, 이를 모방하여 훈련 데이터 분포를 초기에 변형하는 USEFUL 알고리즘을 제안함으로써, 기존 SGD 기반 모델의 성능을 획기적으로 향상시키고 SOTA 를 달성했습니다.