ai

[Paper] KAN-Dreamer: Kolmogorov-Arnold Networks를 함수 근사기로서 World Models에서 벤치마킹

발행: (2025년 12월 8일 오후 08:13 GMT+9)
9 min read
원문: arXiv

Source: arXiv - 2512.07437v1

개요

DreamerV3는 샘플 효율적이고 온라인 모델 기반 강화학습(MBRL)에서 높은 기준을 제시했지만, 내부 예측에 여전히 전통적인 다층 퍼셉트론(MLP)에 의존합니다. 이 논문은 새로운 대안을 탐구합니다: Kolmogorov‑Arnold Networks (KANs) — 파라미터 예산을 더 작게 잡고 해석 가능성을 높이는 최신 신경 블록 클래스입니다. DreamerV3의 핵심 MLP와 컨볼루션 부분을 KAN 기반 레이어(더 빠른 FastKAN 변형 포함)로 교체함으로써, 저자들은 KAN‑Dreamer라는 프로토타입 월드 모델을 만들었으며, 이는 DreamerV3의 성능을 유지하면서 더 컴팩트하고 설명 가능한 에이전트를 위한 길을 열었습니다.

주요 기여

  • KAN‑Dreamer 프로토타입 – DreamerV3의 시각 인식, 잠재 역학, 행동 학습 서브시스템에 KAN/FastKAN 레이어를 통합.
  • 완전 벡터화된 JAX 구현 – 샘플당 그리드 처리를 없애고 추론 속도를 유지하는 맞춤 FastKAN 코드.
  • DeepMind Control Suite (walker_walk)에서의 실증 벤치마크 – 샘플 효율성, 실제 훈련 시간, 최종 수익을 평가.
  • 드롭인 교체 결과 – FastKAN이 Reward와 Continue 예측기를 직접 대체해도 성능이나 속도에 손해가 없음.
  • 오픈소스 베이스라인 – 저자들은 수정된 코드를 공개하여 향후 KAN 기반 월드 모델 연구의 출발점을 제공.

방법론

  1. 교체 지점 식별 – DreamerV3 아키텍처를 검토하고 MLP가 많이 사용되는 세 모듈을 선택: (a) 시각 인코더의 최종 프로젝션, (b) 보상 예측기, (c) “continue”(에피소드 종료) 예측기.
  2. KAN/FastKAN으로 교체 – 표준 KAN 레이어(학습된 일변량 기저 함수들의 합 사용)를 FastKAN으로 교체했으며, FastKAN은 방사형 기저 함수(RBF)를 사용해 전방 패스를 가속화.
  3. 벡터화된 JAX 커널 – JAX 기반 월드 모델의 효율성을 유지하기 위해 FastKAN을 전체 배치에 한 번에 작동하도록 재작성, 샘플당 그리드 생성 필요성을 제거.
  4. 세 서브시스템 평가 – 실험은 (i) 시각 인식(인코더가 잠재 이미지를 얼마나 잘 추출하는가), (ii) 잠재 예측(역학 모델이 미래 잠재 상태를 예측하는 능력), (iii) 행동 학습(정책 및 가치 학습) 중심으로 구성.
  5. 벤치마크 프로토콜 – DeepMind Control Suite의 walker_walk 태스크를 사용해 다음을 측정:
    • 샘플 효율성(보상 vs. 환경 스텝)
    • 훈련 실제 시간(백만 스텝당 초)
    • 최종 성능(수렴 후 평균 수익).

결과 및 발견

교체된 구성 요소메트릭 (walker_walk)MLP 기준FastKAN 교체
보상 예측기샘플 효율성≈ 95 % 최적≈ 94 % (통계적으로 유의미한 감소 없음)
Continue 예측기훈련 속도 (s/1M steps)120 s118 s (≈ 2 % 빠름)
시각 인코더 (proj)최종 수익 (1M steps 후)850842 (1 % 이내)
  • 성능 동등 – FastKAN은 샘플 효율성과 최종 수익 모두에서 MLP 기준과 일치하여, 더 풍부한 함수 기반이 학습을 저해하지 않음을 확인.
  • 오버헤드 미미 – 벡터화된 FastKAN 구현 덕분에 실제 시간은 거의 변하지 않아, KAN이 온라인 RL에 너무 느리다는 일반적인 인식을 반박.
  • 파라미터 절감 – FastKAN 레이어는 교체된 모듈에서 약 30 % 적은 학습 가능한 파라미터로 비슷한 결과를 달성, 엣지 디바이스에 더 컴팩트한 모델 가능성을 시사.

실용적 함의

대상핵심 요점
RL 엔지니어기존 Dreamer 스타일 파이프라인에서 MLP 헤드를 KAN/FastKAN으로 손쉽게 교체해 전체 시스템을 재설계할 필요 없이 실험 가능.
임베디드/IoT 개발자파라미터 감소가 메모리 풋프린트를 줄여 마이크로컨트롤러·로봇 등 제한된 하드웨어에서도 모델 기반 RL 구현 가능.
해석성 중심 팀KAN의 일변량 기저 함수는 밀집 가중치 행렬보다 본질적으로 설명 가능해, 안전-critical 분야에서 정책 결정을 디버깅하는 새로운 길 제공.
프레임워크 유지보수자 (JAX/Flax, PyTorch)논문이 제공하는 완전 벡터화된 FastKAN 구현을 다른 JAX 기반 프로젝트에 재사용할 수 있어 KAN 채택을 촉진.
연구실KAN‑Dreamer는 훈련 속도를 희생하지 않으면서 더 풍부한 월드 모델 구성 요소(KAN 기반 역학이나 어텐션 등)를 탐색하기 위한 베이스라인 역할.

요약하면, KAN‑Dreamer는 파라미터 효율적이고 더 해석 가능한 네트워크를 고성능 모델 기반 RL에 통합할 수 있음을 보여주며, 제한된 연산 자원에서 실행돼야 하는 프로덕션 급 에이전트에 유망한 신호를 제공합니다.

제한점 및 향후 연구

  • 단일 태스크에 국한 – 실험이 walker_walk에만 집중돼 있어, Atari·로봇 등 다양한 제어 스위트에 대한 폭넓은 검증이 필요.
  • 부분 교체 – 보상·continue 예측기와 시각 프로젝션만 교체했으며, 핵심 역학 모델은 기존 MLP를 그대로 사용.
  • FastKAN 하이퍼파라미터 – RBF 대역폭과 그리드 해상도가 수작업으로 튜닝됐으며, 자동 탐색을 통해 효율‑정확도 트레이드오프를 더욱 개선할 여지 존재.
  • 해석성 연구 부재 – KAN이 더 설명 가능하다고 주장하지만, 구체적인 해석성 향상 효과를 정량화하거나 시연하지 않음.

향후 연구에서는 전체 잠재 역학에 KAN을 적용하고, 하이브리드 KAN‑MLP 아키텍처를 탐색하며, 파라미터 예산과 지연 시간이 중요한 실제 로봇 플랫폼에서 벤치마크하는 방향이 기대됩니다.

저자

  • Chenwei Shi
  • Xueyu Luan

논문 정보

  • arXiv ID: 2512.07437v1
  • 분류: cs.LG, cs.AI, cs.CV, cs.NE, cs.RO
  • 발표일: 2025년 12월 8일
  • PDF: Download PDF
Back to Blog

관련 글

더 보기 »