리버스 블록을 활용한 메모리 효율적인 MoE LLM 전체 파라미터 파인튜닝

초록

RevFFN은 Mixture‑of‑Experts(모델) 기반 대형 언어 모델의 전체 파라미터 파인튜닝 시 발생하는 메모리 부담을 크게 낮추기 위해, 입력을 역추적할 수 있는 가역형 Transformer 블록을 설계하였다. 가역성을 이용해 역전파 단계에서 중간 활성값을 재구성함으로써, 대부분의 활성값을 메모리에 저장할 필요가 없어지고, 단일 소비자급 GPU에서도 전체 파라미터 파인튜닝이 가능해진다.

상세 요약

RevFFN 논문은 대규모 언어 모델(LLM) 파인튜닝에서 가장 큰 제약 중 하나인 활성값 저장 메모리 문제를 가역 신경망(reversible neural network) 개념을 통해 근본적으로 해결하고자 한다. 기존의 가역 블록은 입력과 출력 사이에 1:1 매핑을 유지하면서, 역전파 시 입력을 출력만으로 복원할 수 있도록 설계된다. RevFFN은 이러한 가역성을 Transformer 구조에 적용하면서, 특히 MoE(전문가 혼합) 레이어와의 호환성을 확보한다. MoE 레이어는 라우팅 토큰마다 서로 다른 전문가(Feed‑Forward Network)를 선택해 연산량을 효율화하지만, 라우팅 정보와 전문가별 활성값이 동시에 저장돼 메모리 사용량이 급증한다. RevFFN은 (1) 가역형 FFN 블록을 도입해 전문가별 가중치를 그대로 유지하면서도, 전문가 출력값을 역전파 시 재구성하고, (2) 라우팅 결정 자체를 가역 연산으로 처리해 라우팅 매트릭스를 별도 저장하지 않는다. 이때, 라우팅 스코어는 Softmax 전후의 로그 확률 형태로 저장하고, 역전파 시에는 동일한 로그 확률을 재계산해 라우팅 결과를 복원한다.

또한, RevFFN은 메모리 절감 효과를 정량적으로 평가한다. 실험에서는 70B 파라미터 MoE 모델을 기준으로, 기존 DeepSpeed ZeRO‑3 혹은 FSDP와 비교했을 때 피크 메모리 사용량이 45% 이상 감소했으며, GPU당 배치 크기도 2배 이상 늘릴 수 있었다. 중요한 점은 가역 블록이 추가적인 연산 비용을 거의 유발하지 않는다는 것이다. 역전파 시 입력을 재구성하는 과정은 선형 연산과 작은 비선형 연산으로 구성돼, 전체 학습 속도는 기존 방법 대비 5~10% 정도만 감소했다.

RevFFN이 제시하는 핵심 아이디어는 “활성값 저장을 포기하고, 역전파 시 재구성한다”는 전통적인 메모리‑연산 트레이드오프를 뒤집는 것이다. 이는 특히 GPU 메모리가 제한된 환경(예: 24GB 이하)에서 전체 파라미터 파인튜닝을 가능하게 하며, MoE 모델의 전문가 수와 라우팅 복잡성을 유지하면서도 메모리 효율성을 크게 향상시킨다. 향후 연구에서는 가역 블록을 더 깊은 레이어에 적용하거나, 다른 형태의 전문가 라우팅(예: 토큰‑레벨 동적 라우팅)과 결합해 메모리‑연산 효율을 극대화하는 방안을 탐색할 여지가 있다.