보조 손실을 이용한 분리 학습의 통신·메모리 효율화

초록

본 논문은 기존 Split Learning에서 발생하는 양방향 통신과 높은 메모리 사용량을 줄이기 위해, 클라이언트 측에 경량 보조 분류기를 삽입해 로컬 손실을 이용한 독립적인 학습을 제안한다. 서버는 기존의 전역 손실을 사용해 별도로 업데이트하고, 클라이언트와 서버는 한 번의 순방향 전송만으로 학습이 진행된다. CIFAR‑10/100 실험에서 정확도는 기존 방법과 동등하면서 통신량을 약 50%, 메모리 사용량을 최대 58% 절감한다.

상세 분석

이 연구는 Split Learning(SL)의 핵심 병목인 “양방향 전파”와 “전역 활성값 저장”을 근본적으로 재구성한다. 기존 SL은 클라이언트가 중간 활성값 z를 서버에 전송하고, 서버는 손실 L을 계산한 뒤 ∂L/∂z 를 다시 클라이언트로 반환한다. 이 과정은 매 iteration마다 두 차례 네트워크 교환을 요구하고, 클라이언트는 서버로부터 gradient를 받을 때까지 대기해야 하므로 지연(latency)과 대역폭 소모가 크게 늘어난다. 또한, 양쪽 모두 역전파에 필요한 중간 활성값을 메모리에 보관해야 하므로 메모리 부담이 가중된다.

논문은 이러한 구조적 제약을 “보조 손실(auxiliary loss)”이라는 로컬 목표로 대체한다. 클라이언트 측에 얕은 MLP 형태의 보조 분류기 Cₐ를 삽입하고, Cₐ가 예측한 ˜y와 정답 y 사이의 손실 Lₐ를 최소화하도록 학습한다. 이때 클라이언트는 자체적으로 ∂Lₐ/∂θ_b 와 ∂Lₐ/∂θ_a 를 계산해 M_b와 Cₐ를 즉시 업데이트한다. 즉, 서버로부터 어떠한 gradient도 받지 않는다. 서버는 기존과 동일하게 z를 받아 전역 손실 L을 계산하고, ∂L/∂θ_t 로 M_t 를 업데이트한다. 서버는 ∂L/∂z 를 내부적으로만 사용하고 클라이언트에 전송하지 않으므로, 통신은 “클라이언트 → 서버” 한 방향만 존재한다.

핵심 설계 포인트는 두 가지이다. 첫째, 보조 분류기의 용량을 제한해 클라이언트가 과도하게 z에 의존하도록 만들지 않는다. 과도한 보조 모델은 M_b가 서버가 필요로 하는 표현과 다르게 최적화될 위험이 있다. 둘째, 보조 손실에 가중치 λ를 도입해 전역 손실과의 균형을 조정할 수 있다. 논문에서는 λ=0(완전 분리) 설정을 사용해 통신·메모리 절감 효과를 극대화했으며, 필요 시 λ>0을 통해 약한 피드백을 제공할 수 있다.

이러한 설계는 다음과 같은 장점을 제공한다. 1) 통신량 절감: 전통적인 SL이 매 iteration마다 활성값 z와 gradient ∂L/∂z 를 교환하는 반면, 제안 방식은 z만 전송한다. 활성값 크기가 gradient와 동일하거나 더 작기 때문에 이론적으로 50% 이상의 대역폭 절감이 가능하다. 2) 메모리 효율: 클라이언트는 M_b와 Cₐ의 역전파에 필요한 활성값만 보관하고, 전송 후 즉시 해제한다. 서버 역시 ∂L/∂z 를 저장할 필요가 없어 메모리 피크가 감소한다. 3) 동시성: 클라이언트와 서버가 서로 독립적으로 학습을 진행하므로, 서버는 클라이언트가 gradient를 기다리는 동안 다른 배치를 받아 처리할 수 있다. 이는 파이프라인 병렬화와 유사한 효율성을 제공한다.

실험에서는 ResNet‑110을 CIFAR‑10/100에 적용했으며, 세 가지 서로 다른 cut‑layer 위치(얕은, 중간, 깊은)에서 성능을 평가했다. 결과는 전통적인 SL과 거의 동일한 정확도(≤1% 차이)를 보였으며, 통신량은 평균 48%~52% 감소, 클라이언트 메모리 피크는 45%~58% 감소했다. 특히, 보조 분류기의 파라미터 수가 0.5% 수준에 불과함에도 불구하고 학습 안정성이 유지되었으며, λ를 0.1 정도로 조정하면 더 높은 정확도를 얻을 수 있음을 보였다.

전체적으로 이 논문은 “뒤로 전파 없이도 분산 학습이 가능하다”는 새로운 패러다임을 제시한다. 보조 손실을 통한 로컬 최적화는 기존의 전역 gradient 의존성을 대체하면서도 모델 성능을 크게 희생하지 않는다. 향후 연구에서는 보조 손실의 설계(예: contrastive loss, self‑supervised loss)와 다중 클라이언트 환경에서의 동기화 전략을 확장함으로써, 보다 다양한 실세계 엣지 컴퓨팅 시나리오에 적용할 수 있을 것으로 기대된다.