[논문 리뷰] evosax: JAX-based Evolution Strategies

Robert Tjarko Lange|arXiv (Cornell University)|2022. 12. 08.

Evolutionary Algorithms and Applications인용 수 5

한 줄 요약

evosax는 JAX의 JIT 컴파일, 자동 벡터화, 다중 장치 병렬 처리 기능을 활용하여 고성능 하드웨어 가속화된 진화 전략(ES)을 가능하게 하는 JAX 기반 라이브러리입니다. 유한차분, 자연 진화 전략, 유전 알고리즘을 포함한 30종의 ES 알고리즘을 지원하여 단일 명령어로 병렬화할 수 있는 GPU 및 TPU에서 확장 가능한 블랙박스 최적화를 구현합니다.

ABSTRACT

The deep learning revolution has greatly been accelerated by the 'hardware lottery': Recent advances in modern hardware accelerators and compilers paved the way for large-scale batch gradient optimization. Evolutionary optimization, on the other hand, has mainly relied on CPU-parallelism, e.g. using Dask scheduling and distributed multi-host infrastructure. Here we argue that also modern evolutionary computation can significantly benefit from the massive computational throughput provided by GPUs and TPUs. In order to better harness these resources and to enable the next generation of black-box optimization algorithms, we release evosax: A JAX-based library of evolution strategies which allows researchers to leverage powerful function transformations such as just-in-time compilation, automatic vectorization and hardware parallelization. evosax implements 30 evolutionary optimization algorithms including finite-difference-based, estimation-of-distribution evolution strategies and various genetic algorithms. Every single algorithm can directly be executed on hardware accelerators and automatically vectorized or parallelized across devices using a single line of code. It is designed in a modular fashion and allows for flexible usage via a simple ask-evaluate-tell API. We thereby hope to facilitate a new wave of scalable evolutionary optimization algorithms.

연구 동기 및 목표

• 기존에 CPU 기반 병렬 처리에 의존해 온 진화 계산에서 현대적 하드웨어 가속기(GPU/TPU)의 잠재적 활용이 부족한 문제를 해결합니다.
• JAX의 고성능 계산 스택으로 진화 전략을 이식하여 효율적이고 확장 가능한 블랙박스 최적화(BBO)를 가능하게 합니다.
• 하드웨어 병렬 실행, 자동 벡터화, 모듈러한 설계를 통해 차세대 진화 알고리즘의 개발을 촉진합니다.
• 신경망 진화, 메타학습, 아키텍처 탐색 등 다양한 최적화 워크로드를 지원하기 위해 확장 가능한 유틸리티와 인코딩 방식을 제공합니다.
• 기울기 기반 딥러닝과 유도 기반 최적화 간 격차를 해소하기 위해 ES를 현대적 가속기 하드웨어에 네이티브로 호환 가능하게 합니다.

제안 방법

• JAX를 활용해 진화 전략을 구현하여 JIT 컴파일, 자동 미분, 다중 장치에서의 하드웨어 병렬 실행 기능을 활용합니다.
• 어떤 목적 함수(스토케스틱 롤아웃, 신경망 평가 포함)와도 원활하게 통합 가능한 모듈러한 ask-evaluate-tell API를 제공합니다.
• 유한차분 기반 방법(예: OpenAI-ES, ARS), 자연 진화 전략(예: SNES, xNES), CMA-ES 변종을 포함한 총 30종의 독립된 ES 알고리즘을 지원합니다.
• JAX의 `pmap`과 `vmap`를 통한 하드웨어 병렬 실행을 통합하여 단일 코드 라인으로 여러 랜덤 시드 또는 하이퍼파라미터 설정에 대해 병렬로 ES 실행을 수행할 수 있습니다.
• 평탄한 벡터에서 신경망 가중치로의 매개변수 재구성, 피트니스 조정(z-스코어링, 순위 기반), 간접 인코딩(예: 하이퍼네트워크, 랜덤 행렬 투영) 등의 유틸리티를 제공합니다.
• JAX의 집합 연산(예: `pmean`)을 활용한 메모리 최적화 및 장치 간 서브-풀 관리 기능을 통해 대규모 문제에 대한 효율적인 실행을 지원합니다.

Figure 1. Comparison of different evolution strategies (OpenAI-ES, PGPE, SNES and Sep-CMA-ES) implemented by evosax on four continuous control tasks using the Brax library (Freeman et al . , 2021 ) . The simulations use a population size of 256 and the results are averaged over 3 independent runs an

실험 결과

연구 질문

• RQ1 JAX의 기능적 프로그래밍 스택을 활용해 현대적 하드웨어 가속기(GPU/TPU)에서 진화 전략을 효율적으로 가속화할 수 있는가?
• RQ2 JAX의 자동 벡터화 및 하드웨어 병렬 처리 기능이 블랙박스 최적화의 처리량과 확장성에 얼마나 기여하는가?
• RQ3 다양한 최적화 작업에서 JAX 컴파일된 진화 전략은 기존 CPU 기반 구현 대비 성능 및 안정성 면에서 어떻게 비교되는가?
• RQ4 evosax의 모듈러하고 조합 가능한 유틸리티는 메타학습된 진화 전략 및 간접 인코딩 분야에서 새로운 연구 방향을 어떻게 가능하게 하는가?
• RQ5 하드웨어 가속화가 고차원적이고 비미분 가능한 최적화 문제에서 ES의 수렴성과 샘플 효율성에 어떤 영향을 미치는가?

주요 결과

• evosax는 JAX의 JIT 컴파일과 다중 장치 병렬 처리를 통해 진화 전략의 완전한 하드웨어 가속화를 가능하게 하여 대규모 ES 실행의 월클럭 시간을 크게 단축시킵니다.
• 30종의 다양한 진화 전략(유한차분, 자연 진화, CMA-ES 변종 포함)을 지원하며, Brax 환경을 사용한 연속 제어 작업에서 일관된 성능을 보입니다.
• 네 가지 신경진화 작업에 대한 실험 결과, JAX 컴파일된 ES(예: OpenAI-ES, SNES, Sep-CMA-ES)는 최소한의 코드 변경으로 안정적이고 재현 가능한 성능을 달성합니다.
• `pmap`를 통한 단일 명령어 병렬화를 통해 아키텍처 변경 없이도 다양한 랜덤 시드 또는 하이퍼파라미터 설정에서 여러 ES 인스턴스를 동시에 실행할 수 있습니다.
• 피트니스 조정, 매개변수 재구성, 간접 인코딩 유틸리티는 ES를 딥러닝 파ip라인과 복잡한 아키텍처에 쉽게 통합할 수 있도록 크게 단순화합니다.
• 라이브러리의 모듈러한 설계 덕분에 이전에 계산 병목 현상으로 제한되었던 고급 연구 분야, 예를 들어 학습된 진화 전략, 서브-풀 관리, 전략 앙상블 등이 가능해졌습니다.

Figure 2. Comparison of different ES (OpenAI-ES, PGPE, SNES and Sep-CMA-ES) across population sizes on four continuous control tasks. The results are averaged over 3 independent runs.