[논문 리뷰] Single-Shot Pruning for Offline Reinforcement Learning

Samin Yeasar Arnob, Riyasat Ohib외 2명|arXiv (Cornell University)|2021. 12. 31.

Machine Learning and Data Classification인용 수 4

한 줄 요약

이 논문은 오프라인 강화학습(offline RL)에서 훈련 이전에 적용하는 단일 스텝 프루닝 방법을 제안한다. 고정된 데이터셋을 활용하여 희박하고 높은 성능을 보이는 하위신경망을 식별한다. 초기화 시 SNIP 및 GraSP 프루닝을 적용함으로써, 조밀한 네트워크와 비교해도 성능이 유사한데, 이는 95%의 희박성에서도 성능 저하 없이 달성되며, 압축 없이도 메모리 사용량을 4배 줄이고 자원 제약이 있는 장치에서 더 빠른 추론을 가능하게 한다.

ABSTRACT

Deep Reinforcement Learning (RL) is a powerful framework for solving complex real-world problems. Large neural networks employed in the framework are traditionally associated with better generalization capabilities, but their increased size entails the drawbacks of extensive training duration, substantial hardware resources, and longer inference times. One way to tackle this problem is to prune neural networks leaving only the necessary parameters. State-of-the-art concurrent pruning techniques for imposing sparsity perform demonstrably well in applications where data distributions are fixed. However, they have not yet been substantially explored in the context of RL. We close the gap between RL and single-shot pruning techniques and present a general pruning approach to the Offline RL. We leverage a fixed dataset to prune neural networks before the start of RL training. We then run experiments varying the network sparsity level and evaluating the validity of pruning at initialization techniques in continuous control tasks. Our results show that with 95% of the network weights pruned, Offline-RL algorithms can still retain performance in the majority of our experiments. To the best of our knowledge, no prior work utilizing pruning in RL retained performance at such high levels of sparsity. Moreover, pruning at initialization techniques can be easily integrated into any existing Offline-RL algorithms without changing the learning objective.

연구 동기 및 목표

• 오프라인 RL에서 큰 신경망의 높은 계산 및 메모리 비용 문제를 해결하기 위해 훈련 이전에 효율적인 모델 압축을 가능하게 하기 위해.
• 초기 설계된 감독 학습을 위한 단일 스텝 프루닝 기법이 고정된 데이터셋을 가진 오프라인 RL에 효과적으로 적용될 수 있는지 탐색하기 위해.
• 연속 제어 과제에서 극한의 희박성(최대 95%) 조건에서도 프루닝된 네트워크의 성능을 평가하기 위해.
• 초기화 시 프루닝이 기존 오프라인 RL 알고리즘에 손상 없이 원활하게 통합될 수 있음을 입증하기 위해.
• 로봇 공학 및 엣지 AI와 같은 실시간, 자원 제약이 있는 애플리케이션에서 더 빠른 훈련 및 추론을 가능하게 하기 위해.

제안 방법

• 훈련 이전에 무작위로 초기화된 신경망에 단일 스텝 프루닝 기법(SNIP 및 GraSP)을 적용한다.
• D4RL 벤치마크에서 확보한 고정된 데이터셋을 사용하여 기울기 민감도(GraSP) 또는 연결 중요도(SNIP) 기반의 프루닝 기준을 계산한다.
• 재학습이나 반복적 미세조정 없이 초기화 단계에서 한 번만 네트워크를 프루닝하여 희박한 하위신경망을 생성한다.
• 연속 제어 환경(HalfCheetah, Hopper, Walker2D)에서 표준 오프라인 RL 알고리즘(BCQ, BC)을 사용해 프루닝된 네트워크를 평가한다.
• 성능, 메모리 사용량, 레이어 별 희박도를 측정하여 프루닝이 일반화 및 효율성에 미치는 영향을 평가한다.
• PyTorch의 희박 텐서 기능을 사용하여 조밀한 모델과 프루닝된 모델 간의 메모리 요구량을 저장하고 비교한다.

실험 결과

연구 질문

• RQ1 초기화 시 단일 스텝 프루닝이 훈련 이전에 적용되었을 때 오프라인 RL에서 성능을 유지할 수 있는가?
• RQ2 성능 저하 없이 달성 가능한 최대 희박성 수준(예: 95%)은 무엇인가?
• RQ3 다른 프루닝 기준(SNIP 대비 GraSP)이 오프라인 RL에서 프루닝된 네트워크의 구조와 성능에 어떤 영향을 미치는가?
• RQ4 프루닝이 오프라인 RL 모델의 메모리 사용량을 얼마나 줄이고 추론 효율성을 얼마나 향상시키는가?
• RQ5 초기화 시 프루닝이 학습 목표를 수정하지 않고도 기존 오프라인 RL 알고리즘에 통합될 수 있는가?

주요 결과

• 제안된 방법은 95%의 희박성에서도 조밀한 네트워크와 유사한 성능을 유지하며, 여러 환경에서 수익에 유의미한 하락이 없었다.
• GraSP를 사용한 프루닝은 SNIP보다 최종 레이어에 더 많은 가중치를 유지하며, 기울기 흐름을 유지하는 목적에 부합한다.
• 압축 없이도 95% 희박성 모델의 메모리 사용량은 조밀한 모델 대비 4배 감소하였다.
• 이 방법은 즉시 사용 가능한 플러그 앤 플레이 방식이다. 학습 목표를 수정하거나 추가 훈련이 필요 없이 어떤 오프라인 RL 알고리즘에도 적용 가능하다.
• 실험 결과, 배치 크기나 프루닝 반복 수를 변경해도 단일 스텝 프루닝 루프를 초과해 성능 향상이 없음을 확인하였다.
• 결과적으로 고정된 데이터셋을 가진 오프라인 RL 환경이 효과적인 단일 스텝 프루닝을 가능하게 하며, 온라인 RL에서는 데이터 분포의 변화로 인해 이는 불가능하다.