One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning

이 논문은 이기적인 다중 작업 환경에서 그래디언트 충돌을 완화하기 위해 전문가 혼합 (MoE) 아키텍처를 도입하고 학습 과정 중 동적 파라미터 확장 (DPS) 전략을 적용하여, 단일 모델로도 전문 단일 작업 에이전트와 경쟁력 있는 성능을 달성하면서도 환경 상호작용을 71.5% 로 줄인 'ScaleZero'를 제안합니다.

핵심

🎬 1. 문제 상황: "만능 요리사"의 좌절

상상해 보세요. 한 명의 요리사 (AI) 가 있습니다. 이 요리사는 아이스크림을 만드는 일부터 불로 구운 스테이크, 그리고 정교한 디저트까지 모든 요리를 한 번에 배우려고 합니다.

• 기존 방식 (UniZero 등): 이 요리사는 모든 재료를 한 큰 냄비 (단일 모델) 에 넣고 섞어서 끓입니다.
  • 문제점: 아이스크림은 차갑게, 스테이크는 뜨겁게 조리해야 하는데, 같은 냄비에서 동시에 하려니 서로 방해가 됩니다. (이를 **'경쟁하는 그라디언트'**라고 합니다.)
  • 결과: 간단한 요리 (아이스크림) 는 금방 익히지만, 복잡한 요리 (스테이크) 는 자꾸 망칩니다. 요리사의 뇌가 혼란스러워져서 더 이상 새로운 것을 배우지 못하게 됩니다. (이를 **'가소성 붕괴 (Plasticity Collapse)'**라고 합니다.)

🛠️ 2. 해결책 1: "전문가 팀"을 꾸리다 (ScaleZero)

연구자들은 이 문제를 해결하기 위해 요리사의 방식을 완전히 바꿨습니다.

• 새로운 방식 (ScaleZero): 이제 요리사는 혼자 모든 것을 하려 하지 않습니다. 대신 **전문가 팀 (Mixture-of-Experts, MoE)**을 꾸립니다.
  • 비유: 한 명의 요리사가 모든 재료를 섞는 게 아니라, 냉장고 문이 열리면 아이스크림 전문가가, 불이 켜지면 스테이크 전문가가 각각 나와서 자신의 일만 합니다.
  • 효과: 서로 다른 일을 하는 전문가들이 각자만의 공간에서 일하므로, 서로 방해하지 않습니다. 그래서 복잡한 스테이크도 잘 만들고, 간단한 아이스크림도 잘 만들 수 있게 됩니다.
  • 핵심 기술: 이 '전문가 팀' 시스템은 **MoE (Mixture-of-Experts)**라는 기술을 사용했습니다. 입력된 상황에 따라 가장 적합한 전문가만 골라서 일을 시키는 것입니다.

📈 3. 해결책 2: "필요할 때만 인력 추가" (DPS)

하지만 전문가 팀을 처음부터 모두 고용하면 비용이 너무 많이 듭니다. 모든 요리가 동시에 시작되는 것도 아닙니다.

• 새로운 전략 (Dynamic Parameter Scaling, DPS):
  • 비유: 요리사가 처음에는 **기본 조리법 (베이스 모델)**만 배우며 모든 요리를 시도해 봅니다.
  • 진행:
    1. 아이스크림은 금방 잘하게 됩니다. → "이건 끝났으니, 이 부분의 인력은 줄이자." (학습 중단)
    2. 스테이크는 여전히 어렵습니다. → "이건 더 연습이 필요하네! 스테이크 전용 **보조 요리사 (LoRA 어댑터)**를 새로 고용하자!"
  • 효과: 잘하는 일은 더 이상 에너지를 쓰지 않고, 어려운 일에만 집중해서 인력과 시간을 투입합니다.
  • 결과: 같은 실력을 내더라도, 환경과 상호작용하는 횟수 (데이터 수집 비용) 를 약 28.5%나 줄일 수 있었습니다. 마치 "필요한 때에만 필요한 만큼만 투자하는 스마트한 경영"과 같습니다.

🏆 4. 실험 결과: "진짜 만능 선수"가 되다

연구팀은 이 새로운 방식 (ScaleZero) 을 다양한 게임과 환경에서 테스트했습니다.

• 테스트 장소:
  • 아타리 게임 (Atari): 화면을 보고 조이스틱을 조작하는 26 가지 게임 (예: 펭구, 우주선 조종 등).
  • 로봇 제어 (DMC): 물리 법칙을 따르는 로봇 팔이나 다리를 움직이는 18 가지 작업.
  • 텍스트 어드벤처 (Jericho): 글자로 된 미로를 해결하는 게임.
• 결과:
  • **하나의 모델 (ScaleZero)**이 각 게임마다 따로 훈련된 '전문가'들과 거의同等한 (혹은 더 좋은) 실력을 보여주었습니다.
  • 특히 기존 방식이 실패했던 복잡한 게임들에서도 뛰어난 성과를 냈습니다.
  • DPS 전략을 쓰면 더 적은 데이터로도 같은 실력을 낼 수 있었습니다.

💡 5. 요약: 왜 이 연구가 중요한가?

이 논문은 **"하나의 AI 가 모든 일을 동시에 배우는 것"**이 왜 어려운지, 그리고 어떻게 하면 **전문가 팀 (MoE)**과 **스마트한 인력 관리 (DPS)**를 통해 이를 해결할 수 있는지를 증명했습니다.

• 기존의 문제: 모든 일을 한 번에 하려다 보니 뇌가 과부하가 걸려서 망가짐.
• ScaleZero 의 해법:
  1. MoE: 상황에 따라 맞는 전문가만 불러서 일하게 함 (갈등 해결).
  2. DPS: 잘하는 일은 쉬게 하고, 어려운 일에만 추가 인력을 투입함 (효율성 극대화).

결국 이 기술은 **적은 비용으로 더 똑똑하고 다양한 일을 할 수 있는 '일반적인 AI (Generalist Agent)'**를 만드는 데 큰 발걸음이 될 것입니다. 마치 한 명의 천재가 아니라, 상황에 맞춰 최적의 전문가를 불러모으는 훌륭한 프로젝트 매니저를 만든 것과 같습니다.

기술 요약

이 논문은 이질적인 다중 작업 (Heterogeneous Multi-Task) 환경에서 단일 세계 모델 (World Model) 을 사용하여 계획 (Planning) 을 수행할 때 발생하는 근본적인 문제들을 해결하고, 이를 통해 효율적인 범용 에이전트 (Generalist Agent) 를 구축하는 방법을 제시합니다. 제안된 모델은 ScaleZero이며, 학습 과정의 효율성을 높이기 위해 동적 파라미터 스케일링 (Dynamic Parameter Scaling, DPS) 전략을 도입했습니다.

다음은 논문의 기술적 요약입니다.

---

1. 문제 정의 (Problem Definition)

기존의 단일 작업 세계 모델 (예: UniZero) 은 특정 작업에서는 뛰어난 성능을 보이지만, 다양한 작업 (Atari, DMC, 텍스트 게임 등) 을 동시에 학습하는 이질적인 다중 작업 강화학습 (MTRL) 환경에서는 다음과 같은 두 가지 주요 장애물에 직면합니다.

• 가소성 붕괴 (Plasticity Collapse) 및 표현 병목 현상:
  • 단순하고 빠르게 수렴하는 작업의 기울기 (Gradient) 가 복잡한 작업의 기울기를 압도하여, 모델이 복잡한 작업의 특성을 학습하지 못하게 됩니다.
  • 이로 인해 공유된 신경망의 가중치가 새로운 데이터에 적응하는 능력인 **가소성 (Plasticity)**이 상실됩니다.
  • 진단 결과, 복잡한 작업 (예: Seaquest) 에서 성능이 급격히 떨어질 때, 트랜스포머 백본 내의 **휴면 뉴런 비율 (Dormant Neuron Ratio)**이 급증하고 잠재 공간의 노름 (Latent State Norm) 이 통제 없이 팽창하는 현상이 관찰되었습니다. 이는 모델이 새로운 정보를 흡수할 수 있는 능력을 상실했음을 의미합니다.
• 정적 자원 할당 (Static Resource Allocation):
  • 기존 아키텍처는 모든 작업에 대해 동일한 학습 자원 (데이터 수집, 모델 업데이트) 을 균일하게 할당합니다.
  • 이미 학습이 완료된 작업이나 단순한 작업에 불필요한 계산 자원이 소모되고, 여전히 학습이 필요한 난이도 높은 작업에는 자원이 부족해지는 비효율이 발생합니다.

2. 방법론 (Methodology)

저자들은 위 문제를 해결하기 위해 **내부 아키텍처 개선 (ScaleZero)**과 **외부 학습 전략 개선 (DPS)**이라는 두 가지 상호 보완적인 접근법을 제시합니다.

2.1 ScaleZero: 가소성 붕괴를 해결하는 아키텍처

기존 UniZero 의 아키텍처를 분석하고 5 가지 설계 축 (작업 조건부, 인코더, 잠재 정규화, 백본, 최적화) 을 체계적으로 탐색하여 도출된 모델입니다.

• 희소 혼합 전문가 (Sparse Mixture-of-Experts, MoE) 백본:
  • 가장 핵심적인 개선 사항입니다. 기존의 밀집 (Dense) 트랜스포머 피드포워드 네트워크를 MoE로 대체했습니다.
  • 원리: 입력 토큰에 따라 특정 '전문가 (Expert)' 서브네트워크만 활성화되는 조건부 계산 (Conditional Computation) 을 수행합니다.
  • 효과: 각 작업은 서로 다른 전문가 서브네트워크를 통해 처리되므로, 작업 간 **기울기 충돌 (Gradient Conflict)**이 크게 감소합니다. 이는 가소성 붕괴를 방지하고 복잡한 작업에서도 안정적인 학습을 가능하게 합니다.
• ViT 기반 인코더 및 정규화:
  • 시각적 입력 (Atari) 에는 비전 트랜스포머 (ViT) 를, 텍스트 입력 (Jericho) 에는 BGE 인코더를 사용하여 강력한 특징 추출 능력을 확보했습니다.
  • 잠재 상태 정규화에는 SimNorm 대신 표준 LayerNorm 을 사용하여 표현의 유연성을 유지하면서도 안정성을 확보했습니다.

2.2 동적 파라미터 스케일링 (Dynamic Parameter Scaling, DPS)

학습 과정 전체에서 모델 용량을 동적으로 할당하는 온라인 전략입니다.

• 적응형 작업 관리 (Adaptive Task Curation):
  • 사전 정의된 임계값을 만족하여 '해결된 (Solved)' 것으로 판단된 작업은 학습 대상 (Active Set) 에서 제외하여 계산 오버헤드를 제거합니다.
• 단계별 용량 확장 (Staged Capacity Expansion):
  • Stage 0 (Warm-up): 모든 작업에 대해 공유된 베이스 모델 (Base Model) 을 학습합니다.
  • Stage s ≥ 1 (Expansion): 학습 진행 상황에 따라 새로운 LoRA (Low-Rank Adaptation) 어댑터를 점진적으로 주입합니다.
  • 파라미터 격리: 새로운 어댑터를 학습할 때, 이전에 학습된 베이스 모델과 어댑터는 **동결 (Freeze)**시킵니다. 이는 기존 지식의 망각 (Catastrophic Forgetting) 을 방지하고, 새로운 난이도 높은 작업에만 집중된 가소성을 제공합니다.
  • 스케일링 인자: 각 어댑터와 베이스 모델의 기여도를 조절하는 학습 가능한 스케일링 인자 ($\alpha$) 를 도입하여 지식을 재가중치합니다.

3. 주요 기여 (Key Contributions)

1. 가소성 붕괴의 정량적 진단: 통합 세계 모델에서 이질적인 MTRL 환경이 성능 저하와 내부 학습 역학 (휴면 뉴런 증가, 잠재 노름 팽창) 간의 인과 관계를 정량적으로 규명했습니다.
2. ScaleZero 모델 제안: MoE 아키텍처를 통합하여 기울기 충돌을 근본적으로 완화하고, 3 가지 다른 벤치마크 (Atari, DMC, Jericho) 에서 단일 에이전트가 전문 단일 작업 에이전트와 견줄 만한 성능을 발휘함을 입증했습니다.
3. DPS 전략 제안: 학습 진행 상황에 따라 모델 용량을 동적으로 확장하는 LoRA 기반 전략을 제안하여, 환경 상호작용 횟수를 약 28.5% 감소시키면서도 높은 성능을 유지하는 샘플 효율성을 달성했습니다.

4. 실험 결과 (Results)

• Atari 100k 벤치마크 (26 개 게임):
  • ScaleZero(MT, 단일 모델) 는 26 개의 개별적으로 훈련된 UniZero(ST, 전문가 모델) 집합보다 **평균 인간 정규화 점수 (Mean HNS)**가 더 높았습니다.
  • 특히 기존 모델이 실패했던 복잡한 탐색 작업 (Seaquest 등) 에서 뛰어난 성능 개선을 보였습니다.
• DeepMind Control Suite (18 개 연속 제어 작업):
  • ScaleZero 는 대부분의 작업에서 단일 작업 베이스라인과 경쟁력 있는 성능을 보였으며, **중앙값 점수 (Median Score)**가 더 높아 전반적인 견고함을 입증했습니다.
• Jericho (텍스트 기반 어드벤처 게임):
  • 언어 이해와 장기 계획이 필요한 환경에서도 단일 작업 에이전트와 유사한 성능을 달성하여 모델의 모달리티 무관성 (Modality-agnostic) 을 입증했습니다.
• 효율성 평가 (DPS):
  • DPS 전략을 적용한 ScaleZero-DPS 는 DMC 벤치마크에서 목표 성능에 도달하는 데 필요한 환경 상호작용 횟수를 **71.5%**로 줄였습니다 (약 28.5% 절감).

5. 의의 및 결론 (Significance)

이 논문은 단일 모델로 다양한 복잡도와 모달리티를 가진 작업을 동시에 학습하는 '범용 에이전트' 개발에 있어 중요한 진전을 이루었습니다.

• 아키텍처적 통찰: MoE 가 다중 작업 강화학습에서 기울기 충돌을 해결하고 가소성을 유지하는 데 필수적임을 이론적 및 실험적으로 증명했습니다.
• 학습 효율성: 정적 자원 할당의 비효율을 해결하기 위해, 학습 진행도에 따라 모델 용량을 동적으로 조절하는 DPS 전략을 통해 데이터 효율성을 극대화했습니다.
• 미래 방향: 온라인 강화학습의 한계를 극복하기 위해 오프라인 사전 학습 (Offline Pre-training) 과의 결합, 그리고 MoE 와 LoRA 의 더 깊은 시너지 (예: LoRA 를 통한 게이트 네트워크 적응) 를 향후 연구 과제로 제시했습니다.

결론적으로, ScaleZero는 단일 모델로 다양한 작업을 효과적으로 계획하고 학습할 수 있는 강력한 아키텍처적 토대를 제공하며, 효율적인 다중 작업 계획 (Multi-task Planning) 의 새로운 패러다임을 제시합니다.