동적 잔차 연결을 위한 간단한 이중확률 행렬 재구성

초록

본 논문은 기존 Hyper‑Connections(HC)와 Manifold‑Constrained HC(mHC)에서 발생하는 불안정성을 해결하기 위해, Sinkhorn‑Knopp 반복 대신 Birkhoff‑von Neumann 정리를 이용해 퍼뮤테이션 행렬의 볼록 조합으로 이중확률 행렬을 직접 생성하는 mHC‑lite를 제안한다. 이 방법은 정확한 이중확률성을 보장하고, 특수 CUDA 커널 없이도 네이티브 행렬 연산만으로 구현 가능하며, 실험을 통해 mHC‑lite가 mHC와 동등하거나 더 나은 성능을 내면서 학습 속도와 안정성을 크게 향상시킴을 입증한다.

상세 분석

본 논문은 Residual Connection의 확장형인 Hyper‑Connections(HC)가 동적 잔차 행렬 H_res를 통해 여러 스트림 간 정보를 혼합함으로써 표현력을 높이고 수렴 속도를 가속화한다는 점을 출발점으로 삼는다. 그러나 HC는 H_res가 자유롭게 학습되면서 행렬의 스펙트럼 노름이 1을 초과할 위험이 있어, 깊은 네트워크에서 기울기 폭발이나 수렴 불안정을 초래한다. 이를 해결하기 위해 DeepSeek이 제안한 Manifold‑Constrained HC(mHC)는 H_res를 Birkhoff 다면체 B_n(이중확률 행렬 집합) 위에 투사하도록 설계했으며, 이를 구현하기 위해 Sinkhorn‑Knopp(SK) 알고리즘을 20번 반복한다.

하지만 두 가지 근본적인 한계가 존재한다. 첫째, SK는 입력 행렬이 조건이 나쁠 경우 수렴 속도가 매우 느려, 제한된 반복 횟수(예: 20회)만으로는 실제 이중확률성을 만족하지 못한다. 논문은 Linial 등(1998)의 악조건 행렬 예시를 재현해, 20회 SK 후에도 행렬의 열합이 0.92, 0.59 등 크게 벗어나며, 이러한 오차가 층을 거듭할수록 누적돼 24층 네트워크에서는 220%까지 편차가 확대되는 것을 실험적으로 보여준다. 둘째, 효율적인 SK 구현은 반복적인 행/열 정규화를 하나의 fused CUDA 커널로 묶어야 하는데, 이는 특수화된 구현 없이 일반적인 딥러닝 프레임워크에 적용하기 어렵게 만든다.

이러한 문제를 근본적으로 해결하고자 저자들은 Birkhoff‑von Neumann 정리를 활용한다. 정리는 모든 이중확률 행렬이 n!개의 퍼뮤테이션 행렬 P_k의 비음수 가중치 a_k(∑a_k=1)로 표현될 수 있음을 보장한다. 따라서 H_res를 a_k와 고정된 퍼뮤테이션 행렬 집합의 선형 결합으로 직접 파라미터화하면, SK 반복 없이도 정확히 이중확률성을 만족한다. 구현 측면에서는 a_k를 softmax를 통과한 선형 레이어 출력으로 얻고, 미리 준비된 0/1 퍼뮤테이션 행렬 텐서와 행렬 곱을 수행하면 된다. n=4인 경우 n! = 24이므로 가중치 벡터와 퍼뮤테이션 텐서의 차원은 작아 연산 오버헤드가 무시할 수준이며, 기존 mHC와 동일한 구조를 유지하면서 H_res만 교체한다.

실험에서는 ImageNet‑1K, CIFAR‑100, 그리고 대규모 언어 모델(Transformer) 등 다양한 베이스라인에 mHC‑lite를 적용하였다. 결과는 다음과 같다. (1) 정확도 측면에서 mHC‑lite는 mHC와 동등하거나 소폭 상회한다. (2) 학습 시간은 특수 SK 커널이 없는 순수 PyTorch 구현에서도 mHC 대비 1.3~1.7배 빠르며, GPU 메모리 사용량도 감소한다. (3) Gradient norm과 학습 중 발생하는 폭발적 기울기 현상이 mHC‑lite에서는 거의 관찰되지 않아, 이론적으로 보장된 이중확률성 덕분에 안정성이 크게 향상된다. 또한, 1000층 규모의 실험에서도 안정적인 학습이 가능함을 보여, 깊은 네트워크에 대한 확장성을 입증한다.

한계점으로는 퍼뮤테이션 행렬 집합이 n!에 비례해 급격히 커지는 점이다. 현재 실험에서는 n=4(즉, 24개의 퍼뮤테이션)만을 사용했으며, 스트림 수가 크게 늘어날 경우 메모리와 연산 비용이 증가할 수 있다. 저자는 이를 해결하기 위해 샘플링 기반 근사나 저차원 임베딩을 통한 퍼뮤테이션 행렬 압축 등을 향후 연구 과제로 제시한다.

결론적으로, mHC‑lite는 이중확률 행렬을 정확히 구성하는 간단하고 효율적인 방법을 제공함으로써, 기존 mHC가 안고 있던 수렴 불안정성과 구현 복잡성을 동시에 해소한다. 이는 Residual Connection을 다중 스트림으로 확장하려는 모든 최신 아키텍처에 바로 적용 가능한 실용적인 솔루션이다.