[논문 리뷰] Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies

Paul Vicol, Luke Metz외 1명|arXiv (Cornell University)|2021. 12. 27.

Machine Learning and Data Classification인용 수 7

한 줄 요약

이 논문은 히스토리 기반 역전파를 사용하는 비틀림된 역전파를 통해, 비틀림된 계산 그래프에서 편향이 없는 기울기 추정을 위한 지속적 진화 전략(Persistent Evolution Strategies, PES)을 소개한다. PES는 각 부분적 전개 동안의 페르터베이션을 재시작하는 대신 누적시킴으로써, 비틀림 편향을 제거하면서도 낮은 메모리 사용량, 낮은 분산, 빠른 업데이트를 유지한다. 이는 효과적인 초기화수치 최적화 및 학습된 최적화 알고리즘 훈련을 가능하게 한다.

ABSTRACT

Unrolled computation graphs arise in many scenarios, including training RNNs, tuning hyperparameters through unrolled optimization, and training learned optimizers. Current approaches to optimizing parameters in such computation graphs suffer from high variance gradients, bias, slow updates, or large memory usage. We introduce a method called Persistent Evolution Strategies (PES), which divides the computation graph into a series of truncated unrolls, and performs an evolution strategies-based update step after each unroll. PES eliminates bias from these truncations by accumulating correction terms over the entire sequence of unrolls. PES allows for rapid parameter updates, has low memory usage, is unbiased, and has reasonable variance characteristics. We experimentally demonstrate the advantages of PES compared to several other methods for gradient estimation on synthetic tasks, and show its applicability to training learned optimizers and tuning hyperparameters.

연구 동기 및 목표

• 비틀림된 역전파로 인해 발생하는 높은 편향을 해결함으로써, 메타학습 및 초기화수치 최적화에서의 성능 제약을 해소한다.
• TBPTT(편향된 기울기), RTRL(높은 메모리/계산 비용), 표준 ES(비틀림 설정에서 높은 분산) 등의 기존 방법의 한계를 극복한다.
• 편향이 없고 메모리 효율적이며, 순차적 모델에서 실시간, 온라인 파라미터 업데이트에 적합한 기울기 추정 방법을 개발한다.
• 안정적이고 확장 가능한 기울기 추정 메커니즘을 제공함으로써, 학습된 최적화 알고리즘의 효과적인 훈련과 초기화수치 최적화를 가능하게 한다.
• 비틀림된 시스템에서 계산 및 메모리 오버헤드를 줄인, 이론적으로 탄탄하고 실용적인 역전파의 대안을 제공한다.

제안 방법

• PES는 비틀림된 계산 그래프를 짧고 고정된 길이의 세그먼트(비틀림 전개)로 나누고, 각 세그먼트 이후에 진화 전략(ES)을 적용한다.
• 각 전개에서 페르터베이션을 재설정하는 대신, PES는 시간이 지남에 따라 페르터베이션을 누적시킴으로써, 비틀림에도 불구하고 편향 없는 기울기 추정이 가능해진다.
• 기울기 분산을 줄이기 위해 반대칭 샘플링을 사용하여 대칭적인 페르터베이션(양성 및 음성)을 생성함으로써 통계적 효율성을 향상시킨다.
• PES는 모든 입자에 대해 페르터베이션 가중 손실의 평균을 기울기 추정치로 계산하며, 보정 항목은 시간이 지남에 따라 누적된다.
• 최근 전개에서의 분석 기울기를 결합하는 분산 감소 기법을 도입하여 수렴 속도를 향상시킨다.
• JAX의 JIT 컴파일 및 벡터화를 효율적으로 활용하여 높은 병렬성과 낮은 메모리 프로파일을 달성한다.

실험 결과

연구 질문

• RQ1 완전한 역전파의 메모리 및 계산 비용 없이도, 비틀림된 계산 그래프에서 편향 없는 기울기 추정을 달성할 수 있는가?
• RQ2 비틀림된 전개 간에 페르터베이션을 누적시키는 것이 TBPTT와 같은 표준 비틀림된 역전파 방법에서 내재된 편향을 제거하는가?
• RQ3 PES는 순차적 모델에서 실시간, 온라인 파라미터 업데이트를 가능하게 하면서도 낮은 기울기 분산을 유지할 수 있는가?
• RQ4 초기화수치 최적화 및 학습된 최적화 알고리즘 훈련에서, PES는 TBPTT, RTRL, 표준 ES와 비교해 수렴 속도, 안정성, 메모리 사용량 측면에서 어떻게 성능을 내는가?
• RQ5 메타최적화 및 초기화수치 탐색에서 흔히 나타나는 비미분 가능하거나 병리적인 손실 곡면에 대해 PES는 효과적으로 적용될 수 있는가?

주요 결과

• PES는 손실의 스무스한 변형에 대해 기울기 추정치를 편향 없이 제공하며, 이는 이차 손실에 대해 정확히 편향 없이 작동한다.
• 특히 최근 전개에서의 분석 기울기를 결합한 경우, 표준 진화 전략보다 비틀림 설정에서 훨씬 낮은 기울기 분산을 달성한다.
• 초기화수치 최적화 작업에서는 TBPTT와 표준 ES보다 더 빠르고 안정적으로 수렴하며, 최종 성능은 유사하거나 더 우수하다.
• PES는 학습된 최적화 알고리즘의 효과적인 메타학습을 가능하게 하여 내부 루프의 초기화수치 파라미터를 안정적이고 효율적으로 최적화함을 보여준다.
• 지속적 제어 과제에 대해 잘 스케일링되며, 비틀림 최적화 목표를 가진 정책 학습에서 뛰어난 성능을 보인다.
• 실험 결과는 PES가 낮은 메모리 사용량을 유지하면서도 빈번한 파라미터 업데이트를 가능하게 하여 실시간 및 온라인 학습 시나리오에 적합하다는 것을 확인한다.