[논문 리뷰] Learning to Learn with Generative Models of Neural Network Checkpoints

William Peebles, Ilija Radosavovic외 3명|arXiv (Cornell University)|2022. 09. 26.

Model Reduction and Neural Networks인용 수 7

한 줄 요약

이 논문은 2300만 개의 신경망 체크포인트로 구성된 대규모 데이터셋을 훈련하여 신경망 최적화를 학습하는 생성 모델 G.pt를 제안한다. 조건부 디퓨전 트랜스포머를 사용하여, 백프로파게이션 또는 누적 최적화 없이도, 새로운 아키텍처와 작업에 대해 손실, 오차 또는 수익과 같은 목표 지표를 달성하는 데 필요한 최적의 매개변수 업데이트를 한 번의 단계로 예측한다.

ABSTRACT

We explore a data-driven approach for learning to optimize neural networks. We construct a dataset of neural network checkpoints and train a generative model on the parameters. In particular, our model is a conditional diffusion transformer that, given an initial input parameter vector and a prompted loss, error, or return, predicts the distribution over parameter updates that achieve the desired metric. At test time, it can optimize neural networks with unseen parameters for downstream tasks in just one update. We find that our approach successfully generates parameters for a wide range of loss prompts. Moreover, it can sample multimodal parameter solutions and has favorable scaling properties. We apply our method to different neural network architectures and tasks in supervised and reinforcement learning.

연구 동기 및 목표

• 기존 경량 최적화 방법이 과거 경험에서 개선되지 못하는 한계를 해결한다.
• 누적 최적화 또는 강화학습을 통한 학습 최적화 방법의 불안정성과 복잡성을 극복한다.
• 기존의 방대한 양의 신경망 체크포인트를 메타학습 최적화 전략의 데이터 소스로 활용한다.
• 다양한 아키텍처와 작업 전반에서 확장 가능한, 데이터 기반의 한 번의 단계 최적화 방법을 개발한다.
• 분류 오차나 강화학습 수익과 같은 비미분 가능 목적함수의 최적화를 가능하게 한다.

제안 방법

• 지도학습(MNIST, CIFAR-10) 및 강화학습(CartPole) 작업을 포함한 10만 개 이상의 훈련 런에서 유래한 2300만 개의 신경망 체크포인트 데이터셋을 구축한다.
• 표준 최적화 방법(예: Adam, SGD)을 사용하여 매개변수와 작업 수준의 지표(테스트 손실, 오차, 수익)를 포함한 체크포인트를 수집한다.
• 매개변수 수준의 데이터 증강(예: 가중치 순서 변경)을 적용하여 데이터셋의 다양성을 증가시키면서도 모델의 기능을 유지한다.
• 초기 매개변수 벡터와 목표 지표 프롬프트(예: 목표 손실)를 입력으로 받아 최적의 매개변수 업데이트에 대한 분포를 예측하는 조건부 디퓨전 트랜스포머 모델(G.pt)을 훈련한다.
• 누적 최적화나 강화학습을 대체로 표준 생성 모델링 기법을 사용하여 안정적이고 확장 가능한 훈련을 가능하게 한다.
• 추론 시, G.pt의 단일 전방향 패스를 통해 새로운 초기화 및 알려지지 않은 아키텍처에 대해 최적화된 매개변수를 생성한다.

실험 결과

연구 질문

• RQ1 대규모 신경망 체크포인트 데이터셋을 기반으로 훈련된 생성 모델은 한 번의 단계로 효과적인 매개변수 업데이트를 학습할 수 있는가?
• RQ2 이러한 모델은 분포 외의 가중치 초기화 방법과 알려지지 않은 아키텍처로 일반화할 수 있는가?
• RQ3 동일한 최적화 프롬프트에 대해 다중 모달 솔루션을 생성할 수 있는가? 이는 동일한 지표에 도달하는 다양한 경로를 학습했음을 의미한다.
• RQ4 분류 오차나 강화학습 수익과 같은 비미분 가능 목적함수에 대해 이 방법이 효과적으로 작동하는가?
• RQ5 다양한 네트워크 아키텍처와 작업 전반에서 학습된 최적화 전략의 성능는 어떻게 스케일링되는가?

주요 결과

• G.pt는 다양한 아키텍처와 작업에 대해 새로운 초기화 상태에서도 단일 매개변수 업데이트로 신경망 최적화를 성공적으로 수행한다.
• 모델은 분포 외의 가중치 초기화 방법으로도 일반화되어 있으며, 훈련 분포를 초월한 강건성을 보여준다.
• G.pt는 동일한 최적화 프롬프트에 대해 다중 모달 솔루션을 생성할 수 있으며, 이는 최적 성능에 도달하는 다양한 경로를 학습했음을 나타낸다.
• 모델는 다양한 네트워크 깊이와 너비에서 우수한 스케일링 성질을 유지하며 성능을 유지한다.
• 전통적인 학습 최적화 방법과 달리, G.pt는 분류 오차나 강화학습 수익과 같은 비미분 가능 목적함수를 효과적으로 최적화한다.
• CIFAR-10과 같은 복잡한 데이터셋에서는 다소의 과소적합이 관찰되지만, 단순한 작업에서는 뛰어난 성능을 보이며 새로운 아키텍처로의 일반화 능력도 입증된다.