최적 그래디언트 체크포인트 탐색을 통한 임의 계산 그래프 메모리 절감

초록

본 논문은 임의 형태의 계산 그래프(ACG)에서도 최적의 그래디언트 체크포인트(GC)를 자동으로 선택하는 이론과 알고리즘을 제시한다. 선형 계산 그래프(LCG)에서의 기존 휴리스틱 방법을 일반화하여, 메모리 사용량을 최대 80%까지 감소시키면서도 시간 오버헤드는 30~50% 수준으로 제한한다.

상세 분석

이 연구는 딥러닝 훈련 시 GPU 메모리 한계에 직면한 문제를 근본적으로 해결하고자 한다. 기존의 Gradient Checkpointing(GCP)은 일부 텐서만 저장하고, 역전파 단계에서 누락된 텐서를 재전파(local re‑forward)함으로써 메모리를 절감한다. 그러나 어느 텐서를 체크포인트로 선택할지는 모델마다 다르며, 특히 비선형 구조를 갖는 Arbitrary Computation Graph(ACG)에서는 수작업이나 단순 히스토리 기반 휴리스틱이 비효율적이다.

논문은 먼저 메모리 비용 함수를 l(v)로 정의하고, 체크포인트 집합 V_R에 대해 전체 메모리 피크를
 min_{V_R} ( Σ_i l(v_Ri) + max_i l(v_Ri, v_R{i+1}) )
로 수식화한다. 여기서 첫 항은 체크포인트 자체의 메모리, 두 번째 항은 가장 큰 재전파 구간의 메모리를 의미한다.

선형 그래프(LGC)에서는 두 가지 경우를 분석한다. (1) 모든 노드 비용이 동일할 때, 최적 해는 √N개의 체크포인트를 균등하게 배치해 전체 메모리를 2√N으로 감소시키는 것이며, 이는 기존 Chen 알고리즘의 근거가 된다. (2) 비용이 비동일한 경우, 접근성 그래프(Accessibility Graph)를 구축한다. 두 노드 사이의 누적 비용이 제한 C 이하이면 해당 간선을 추가하고, 이 그래프에서 소스‑타깃 간 최단 경로를 찾음으로써 주어진 C 하에서 최소 체크포인트 비용을 구한다. 모든 가능한 C(모든 노드 쌍의 비용) 를 순회하면 원 문제의 전역 최적 해를 얻을 수 있다. 알고리즘 복잡도는 O(|V|^2|E| + |V|^3 log|V|) 로, 실용적인 규모의 네트워크에서도 적용 가능하다.

임의 그래프(ACG)로 확장하기 위해 독립 구간(Independent Segment, IS) 개념을 도입한다. IS는 선택된 두 체크포인트 사이의 서브그래프가 외부와 연결되지 않는 부분 그래프이며, 각 IS는 독립적으로 재전파·역전파가 가능하다. IS를 다시 선형 분할 정점(linear splitting vertex)이나 비선형 분할 정점에 따라 분류하고, 각각에 맞는 분할 전략을 적용한다. 큰 IS가 제한 C를 초과하면 추가 체크포인트를 삽입해 더 작은 IS로 나누는 과정을 반복한다. 최종적으로는 모든 IS의 최대 메모리 비용이 C 이하가 되도록 체크포인트 집합을 최적화한다.

실험에서는 AlexNet, VGG, ResNet, DenseNet, Inception, 그리고 NAS로 탐색된 복잡한 네트워크에 대해 기존 LCG 전용 방법(Griewank‑Walther, Chen, Gruslys)과 비교하였다. 결과는 메모리 절감율이 60~80%에 달하고, 추가 연산 비용은 평균 35% 수준으로, 특히 대규모 입력(예: 600×600 이미지, 64프레임 비디오)에서 기존 GPU 메모리 한계를 크게 완화한다. 코드가 공개되어 재현 가능성이 높으며, 프레임워크(Pytorch, TensorFlow)와도 쉽게 통합될 수 있다.

핵심 기여는 (1) ACG에 대한 최적 GC 선택 이론을 정립, (2) 접근성 그래프와 IS 기반 분할 알고리즘을 통한 효율적 구현, (3) 다양한 실제 모델에 대한 광범위한 실증 검증이다. 제한점으로는 그래프가 매우 큰 경우 접근성 그래프 구축 비용이 증가할 수 있으며, 동적 메모리 할당이 빈번한 모델에서는 추가 오버헤드가 발생할 가능성이 있다. 향후 연구에서는 근사적 그래프 압축 기법과 하드웨어 친화적 스케줄링을 결합해 실행 시간 최적화를 목표로 할 수 있다.