하이퍼 커넥션과 mHC 해설: 트랜스포머 잔차 연결을 학습 가능하게 바꾼 방법과 안정화 기법

핵심 메시지

• 잔차 연결은 깊은 층에서 얕은 층으로 기울기가 곧바로 흐르는 항등 경로를 만들어 깊은 모델의 학습을 가능하게 했다.
• 하이퍼 커넥션은 잔차 경로를 여러 갈래로 늘리고 그 결합 세기까지 학습해, 표준 잔차 연결보다 최대 1.8배 빠른 수렴을 보였다.
• 경로 사이를 섞는 행렬이 제약 없이 학습되면 층을 거치며 값이 폭발하거나 소멸해 학습이 불안정해진다.
• 해법은 그 행렬을 모든 원소가 양수이고 행과 열의 합이 1인 이중 확률 행렬로 만드는 것이며, 지수 함수와 반복 정규화로 간단히 구현된다.
• 확장된 잔차 경로는 메모리를 많이 쓰지만, 커널 융합과 활성값 재계산 등 세 가지 인프라 설계로 확장률 4에서 학습 오버헤드를 6.7%까지 낮췄다.

쉽게 이해하기

층을 깊이 쌓으면 표현력이 커진다는 것이 신경망의 직관이지만, 실제로는 깊은 모델이 훈련 오차조차 더 높게 나오는 역설이 오래 관찰됐다. 얕은 신경망의 층을 복제하고 나머지 층을 항등 사상으로 채우면 최소한 얕은 모델만큼은 나와야 하므로, 해가 없는 것이 아니라 최적화가 어려웠던 것이다. 잔차 학습은 층이 원하는 사상을 직접 배우게 하는 대신 입력과 출력의 차이만 배우게 해 이 문제를 풀었다.

잔차 연결의 힘은 식을 재귀적으로 펼쳐 보면 분명해진다. 깊은 층의 특징은 얕은 층의 특징에 잔차 함수들의 합을 더한 형태가 되고, 이 항등 사상 덕분에 역전파에서 기울기가 중간 가중치 층을 거치지 않고 직접 전달되는 경로가 생긴다. 기울기 소실이 완화되면서 이 설계는 이미지, 언어, 바둑, 단백질 구조 예측까지 사실상 모든 영역으로 퍼졌다.

다만 잔차 경로가 단 하나라는 점은 제약이다. 하이퍼 커넥션은 입력 특징을 n배로 늘려 여러 갈래의 잔차 흐름을 만든 뒤, 층에 넣기 전에 가중합으로 모으고 층의 출력은 학습 가능한 계수로 다시 펼쳐 각 흐름에 더한다. 여기에 흐름 사이를 잇는 학습 가능한 선형 변환을 두면 특징을 어느 흐름에서 어느 흐름으로 보낼지 결정하는 라우터가 된다. 이 계수들은 입력에 따라 달라지도록 정적 항과 동적 항으로 나뉘어 파라미터화된다.

문제는 이 기법을 대규모 모델 학습에 적용하자 심각한 불안정이 나타났다는 점이다. 원인은 구조에 있다. 표준 잔차 연결에서는 얕은 층과 깊은 층 사이에 항등 사상이 놓이지만, 하이퍼 커넥션에서는 그 자리에 층마다의 혼합 행렬들이 연달아 곱해진 합성 사상이 놓인다. 1차원으로 단순화해 보면 계수가 1.1이기만 해도 층을 거치며 값이 100배로 커지고, 1보다 작으면 급격히 사라지며, 음수면 부호가 요동친다.

해결책은 그 혼합 행렬을 얌전한 형태로 강제하는 것이다. 모든 원소가 양수이고 행과 열의 합이 각각 1인 이중 확률 행렬로 만들면 여러 층을 지난 합성 사상도 같은 성질을 유지해 폭발과 소실이 함께 막힌다. 구현은 의외로 단순하다. 원소마다 지수 함수를 씌워 양수로 만든 뒤 행과 열의 합을 1로 번갈아 재조정하는 반복을 몇 번 돌리면 된다. 논문은 이를 이중 확률 행렬의 다양체로 사영한다고 표현한다.

주요 인사이트

• 나머지 두 계수 행렬은 활성 함수를 탄젠트 하이퍼볼릭에서 시그모이드로 바꿔 다듬었다. 음수 계수를 없애고 값의 범위를 묶어 두기 위해서다.
• 층 출력에 곱해지는 계수 앞에 2가 붙는 데는 이유가 있다. 초기화 시 시그모이드 입력이 0에 가까워 출력이 0.5 근처가 되므로, 2를 곱하면 학습 초기에 하이퍼 커넥션이 표준 잔차 연결과 정확히 같게 동작한다.
• 잔차 흐름을 n배로 늘리면 GPU 메모리 사용량과 메모리 접근이 함께 늘어난다. 정규화 순서 재배치와 연산 융합, 순전파 후 중간 활성값을 버리고 역전파 때 다시 계산하는 기법, 계산과 통신을 겹치는 스케줄링으로 이 비용을 눌렀다.
• 표준 잔차 연결은 오랫동안 손대지 않은 구성 요소였다. 이 영역에서 실질적인 개선이 나왔다는 사실 자체가, 이미 굳어진 것처럼 보이는 설계에도 여지가 남아 있음을 보여 준다.

자주 묻는 질문

관련 AI 소식