style="width:100.0%" />

서론

ResNets의 도입 이후, 딥 뉴럴 네트워크 아키텍처는 급속한 진화를 거쳤습니다. Fig. 1(a)에 제시된 것처럼, 단일 계층의 구조는 다음과 같이 표현됩니다:

\begin{equation}
    \mathbf{x}_{l+1} = \mathbf{x}_l + \mathcal{F}(\mathbf{x}_l, \mathcal{W}_l),
    \label{eqn:single_rc}
\end{equation}

여기서 $`\mathbf{x}_l`$와 $`\mathbf{x}_{l+1}`$는 각각 $`C`$-차원 입력과 출력을 나타내며, $`\mathcal{F}`$는 잔차 함수를 의미합니다. 지난 10년 동안 잔차 함수 $`\mathcal{F}`$는 합성곱, 어텐션 메커니즘, 및 피드 포워드 네트워크와 같은 다양한 연산을 포함하면서 진화해왔지만, 잔차 연결의 패러다임은 원래 형태를 유지하고 있습니다. Transformer 아키텍처의 진행과 함께 이 패러다임은 현재 대형 언어 모델(LLMs)에서 기본 설계 요소로 자리잡았습니다.

이 성공은 주로 잔차 연결의 간결한 형식에 귀인됩니다. 무엇보다 중요한 것은 초기 연구가 잔차 연결의 항등 매핑 속성이 대규모 훈련 중 안정성과 효율성을 유지하는 데 기여한다는 것을 밝혀냈습니다. 잔차 연결을 여러 계층으로 재귀적으로 확장하면 Eq. [eqn:single_rc]은 다음과 같이 표현됩니다:

\begin{equation}
    \mathbf{x}_L = \mathbf{x}_l + \sum_{i=l}^{L-1} \mathcal{F}(\mathbf{x}_i, \mathcal{W}_i),
    \label{eqn:multi_rc}
\end{equation}

여기서 $`L`$과 $`l`$은 각각 더 깊고 얕은 계층을 나타냅니다. 항등 매핑이라는 용어는 $`\mathbf{x}_l`$ 자체를 지칭하며, 이는 신호가 수정 없이 얕은 계층에서 더 깊은 계층으로 직접 맵핑된다는 속성을 강조합니다.

최근에는 Hyper-Connections(HC)와 같은 연구들이 잔차 연결에 새로운 차원을 추가하고 그 성능 잠재력을 경험적으로 입증했습니다. HC의 단일 계층 구조는 Fig. 1(b)에서 설명되어 있습니다. HC는 잔차 스트림의 너비를 확장하고 연결 복잡성을 높여서 계산 단위별 FLOPs에 대한 연산 부담을 변경하지 않고도 위상학적 복잡성을 크게 증가시킵니다. 공식적으로, HC의 단일 계층 전파는 다음과 같이 정의됩니다:

\begin{equation}
    \mathbf{x}_{l+1} = \mathcal{H}_{l}^{\mathrm{res}}\mathbf{x}_l + \mathcal{H}_{l}^{\mathrm{post}\, \top}\mathcal{F}(\mathcal{H}_{l}^{\mathrm{pre}}\mathbf{x}_l, \mathcal{W}_l),
    \label{eqn:single_hc}
\end{equation}

여기서 $`\mathbf{x}_{l}`$와 $`\mathbf{x}_{l+1}`$는 각각 $`l`$번째 계층의 입력과 출력을 나타냅니다. Eq. [eqn:single_rc]에서와 달리, $`\mathbf{x}_{l}`$와 $`\mathbf{x}_{l+1}`$의 특성 차원은 $`C`$에서 $`n \times C`$로 확장됩니다. 여기서 $`n`$은 확장 비율을 나타냅니다. $`\mathcal{H}_{l}^{\mathrm{res}} \in \mathbb{R}^{n \times n}`$는 잔차 스트림 내의 특성들을 혼합하는 학습 가능한 매핑을 나타내며, 또한 학습 가능한 매핑인 $`\mathcal{H}_{l}^{\mathrm{pre}} \in \mathbb{R}^{1 \times n}`$은 $`nC`$-차원 스트림에서 특성을 집계하여 $`C`$-차원 계층 입력으로 변환하고, 반대로 $`\mathcal{H}_{l}^{\mathrm{post}} \in \mathbb{R}^{1 \times n}`$은 계층 출력을 다시 스트림에 매핑합니다.

그러나 훈련 규모가 증가함에 따라 HC는 잠재적인 안정성 문제를 유발할 수 있습니다. 주된 우려 사항은 HC의 제약 없는 성격이 여러 계층으로 아키텍처가 확장될 때 항등 매핑 속성을 해치는 것입니다. 여러 병렬 스트림을 포함하는 아키텍처에서 이상적인 항등 매핑은 보존 메커니즘으로 작용하며, 양방향 전파 중 스트림 간의 평균 신호 강도가 불변하도록 합니다. Eq. [eqn:single_hc]를 통해 HC를 여러 계층으로 재귀적으로 확장하면 다음과 같습니다:

\begin{equation}
    \mathbf{x}_{L} = \left(\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}\right)\mathbf{x}_l + \sum_{i=l}^{L-1}\left(\prod_{j=1}^{L-1-i}\mathcal{H}_{L-j}^{\mathrm{res}}\right)\mathcal{H}_{i}^{\mathrm{post}\, \top}\mathcal{F}(\mathcal{H}_{i}^{\mathrm{pre}}\mathbf{x}_i, \mathcal{W}_i),
    \label{eqn:multi_hc}
\end{equation}

여기서 $`L`$과 $`l`$은 각각 더 깊고 얕은 계층을 나타냅니다. Eq. [eqn:multi_rc]와 달리, HC의 복합 매핑 $`\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}`$는 특성들의 전역 평균을 유지하지 못합니다. 이 차이점은 신호가 무한대로 증폭되거나 감소하는 것을 초래하며, 대규모 훈련 중 안정성을 저해합니다. 또한 HC는 FLOPs에 대한 연산 효율성을 유지하나, 확장된 잔차 스트림에 대한 메모리 액세스 비용을 최적화하지 않았습니다. 이러한 요인들은 합계적으로 HC의 실질적인 확장성과 대규모 훈련에서의 활용을 방해합니다.

이러한 도전 과제를 해결하기 위해, 우리는 Fig. 1(c)에 제시된 것처럼 Manifold-Constrained Hyper-Connections (mHC)을 제안합니다. mHC는 HC의 잔차 연결 공간을 특정 다양체로 투영하여 항등 매핑 속성을 복원하면서 철저한 인프라 구조 최적화를 통한 효율성을 보장하는 일반적인 프레임워크입니다. 특히, mHC는 Sinkhorn-Knopp 알고리즘을 활용해 $`\mathcal{H}_{l}^{\mathrm{res}}`$를 Birkhoff 다각형으로 엔트로피적으로 투영합니다. 이 연산은 잔차 연결 행렬들을 쌍대 확률적 행렬들로 구성된 다양체 내에서 제약을 가합니다. 이러한 행렬들의 각 행과 열 합이 $`1`$에 등한하기 때문에, $`\mathcal{H}_{l}^{\mathrm{res}}\mathbf{x}_l`$ 연산은 입력 특성의 볼록 결합으로 기능합니다. 이 특징은 평균을 보존하고 신호 노름을 엄격하게 규제하는 잘 조건화된 신호 전파를 가능케 하며, 사라짐 또는 터지는 신호 위험을 효과적으로 완화시킵니다. 또한 쌍대 확률적 행렬의 행렬 곱셈에 대한 폐합성 덕분에 복합 매핑 $`\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}`$도 이 보존 속성을 유지합니다. 결과적으로 mHC는 임의의 깊이 사이에서 항등 매핑의 안정성을 효과적으로 유지합니다. 효율성을 위해 우리는 커널 융합을 활용하고 TileLang을 사용하여 혼합 정밀도 커널을 개발했습니다. 또한, DualPipe 스케줄 내에서 선택적 재계산과 통신 중복으로 메모리 발자국을 줄였습니다.

언어 모델 사전 훈련에 대한 광범위한 실험은 mHC가 안정성과 확장성을 유지하면서 HC의 성능 이점을 보유하고 있음을 증명했습니다. 내부 대규모 훈련 결과, mHC는 $`n=4`$일 때 6.7%의 추가 시간 부담을 초래하며도 규모에 따른 훈련을 지원합니다.

관련 연구

딥러닝의 아키텍처 발전은 주로 마이크로 설계와 매크로 설계로 분류됩니다. 마이크로 설계는 계산 블록 내부 구조를 지정하며, 공간, 시간 및 채널 차원을 통해 특성이 어떻게 처리되는지 결정합니다. 반면에 매크로 설계는 블록 간의 위상 구조를 설정하여 서로 다른 계층에서 특성 표현이 어떻게 전파되고 결합되는지를 지정합니다.

마이크로 설계

특징 신호의 구조화된 처리에서는 매개변수 공유 및 변환 불변성을 통해 합성곱이 주도했습니다. 이후, 효율성을 최적화하기 위한 다양한 변형들인 depthwise separable 및 그룹 합성곱이 등장했지만, Transformer의 도입은 어텐션과 피드포워드 네트워크(FFNs)를 현대 아키텍처의 기본 구성 요소로 만들었습니다. 어텐션 메커니즘은 전역 정보 전파를 촉진하며, FFN은 개별 특성의 표현력 강화에 기여합니다. LLMs의 성능과 계산 요구 사항 사이의 균형을 위해, 어텐션 메커니즘은 Multi-Query Attention(MQA), Grouped-Query Attention(GQA), 및 Multi-Head Latent Attention(MLA)와 같은 효율적인 변형으로 진화했습니다. 동시에, FFN은 Mixture-of-Experts(MoE)를 통해 희소 계산 패러다임으로 일반화되어 대규모 매개변수 확장에도 불구하고 비례하는 연산 비용 없이 가능해졌습니다.

매크로 설계

매크로 설계는 네트워크의 전역 위상 구조를 규제합니다. ResNet 이후, DenseNet 및 FractalNet은 밀도 연결과 다중 경로 구조를 통해 성능을 향상시키려고 노력했습니다. Deep Layer Aggregation(DLA)은 다양한 깊이와 해상도에서의 특성 집계를 재귀적으로 확장하는 방식으로 이 패러다임을 더욱 발전시켰습니다.

최근 매크로 설계는 잔차 스트림 너비의 확장을 중심으로 연구가 이루어지고 있습니다. Hyper-Connections(HC)은 다양한 깊이에서 특성 간 연결 강도를 조절하기 위한 학습 가능한 행렬을 도입했으며, Residual Matrix Transformer(RMT)은 표준 잔차 스트림 대신 외적 메모리 매트릭스를 사용하여 특성을 저장하는 방식으로 접근했습니다. 또한 MUDDFormer는 다중 경로 동적인 연결을 활용해 계층 간 정보 흐름을 최적화합니다. 그러나 이러한 방법들은 잔차 연결의 본래 항등 매핑 속성을 해치며 안정성과 확장성을 저해하는 한편, 확장된 특성 너비에 따른 메모리 액세스 비용도 증가시킵니다. HC를 기반으로 제안된 mHC는 잔차 연결 공간을 특정 다양체로 투영하여 항등 매핑 속성을 복원하며, 철저한 인프라 구조 최적화를 통해 효율성을 보장합니다. 이 접근법은 확장성과 안정성을 개선하면서도 확장된 연결의 위상학적 이점을 유지합니다.

서두

우리는 이 연구에서 사용할 표기법을 먼저 정립합니다. HC 공식에서, $`l`$번째 계층에 대한 입력 $`\textbf{x}_l\in \mathbb{R}^{1\times C}`$는 확장 비율 $`n`$으로 확장되어 은닉 행렬 $`\textbf{x}_l = (\textbf{x}^\top_{l,0}, \ldots, \textbf{x}^\top_{l,n-1})^\top \in \mathbb{R}^{n \times C}`$을 구성하며 이는 $`n`$ 스트림 잔차로 볼 수 있습니다. 이를 통해 잔차 스트림의 너비가 확장됩니다. 이 스트림의 읽기, 쓰기 및 업데이트 과정을 관리하기 위해 HC는 세 가지 학습 가능한 선형 매핑—$`\ensuremath{\mathcal{H}^{\mathrm{pre}}_{l}}, \ensuremath{\mathcal{H}^{\mathrm{post}}_{l}}\in \mathbb{R}^{1\times n}`$, 그리고 $`\ensuremath{\mathcal{H}^{\mathrm{res}}_{l}}\in \mathbb{R}^{n\times n}`$를 도입합니다. 이러한 매핑은 표준 잔차 연결인 Eq. [eqn:single_rc]을 수정하여 Eq. [eqn:single_hc]의 공식으로 이끕니다.

HC에서 학습 가능한 매핑은 두 가지 계수 부분으로 구성됩니다: 입력에 의존적인 동적 매핑과 전역적인 정적 매핑입니다. 공식적으로, HC는 다음과 같이 계수를 계산합니다:

\begin{equation}
    \begin{cases}
        \tilde{\mathbf{x}}_l = \text{RMSNorm}(\mathbf{x}_l) \\
        \ensuremath{\mathcal{H}^{\mathrm{pre}}_{l}} = \alpha_l^\mathrm{pre} \cdot \tanh(\theta^\mathrm{pre}_l \tilde{\mathbf{x}}^\top_l) + \mathbf{b}_l^\mathrm{pre} \\
        \ensuremath{\mathcal{H}^{\mathrm{post}}_{l}} = \alpha_l^\mathrm{post} \cdot \tanh(\theta^\mathrm{post}_l \tilde{\mathbf{x}}^\top_l) + \mathbf{b}_l^\mathrm{post} \\
        \ensuremath{\mathcal{H}^{\mathrm{res}}_{l}} = \alpha_l^\mathrm{res} \cdot \tanh(\theta^\mathrm{res}_l \tilde{\mathbf{x}}^\top_l) + \mathbf{b}_l^\mathrm{res}, \\
    \end{cases}
    \label{eqn:hc_details}
\end{equation}

여기서 $`\text{RMSNorm}(\cdot)`$은 마지막 차원에 적용되며, 스칼라 $`\alpha_l^\mathrm{pre}, \alpha_l^\mathrm{post}`$와 $`\alpha_l^\mathrm{res} \in \mathbb{R}`$는 초기값이 작은 학습 가능한 게이팅 요소입니다. 동적 매핑은 $`\theta^\mathrm{pre}_l, \theta^\mathrm{post}_l \in \mathbb{R}^{1 \times C}`$와 $`\theta^\mathrm{res}_l \in \mathbb{R}^{n \times C}`$로 파라미터화된 선형 프로젝션을 통해 도출되며, 정적 매핑은 학습 가능한 편향들인 $`\mathbf{b}_l^\mathrm{pre}, \mathbf{b}_l^\mathrm{post} \in \mathbb{R}^{1\times n}`$와 $`\mathbf{b}_l^\mathrm{res} \in \mathbb{R}^{n\times n}`$로 표현됩니다.

이러한 매핑들—$`\ensuremath{\mathcal{H}^{\mathrm{pre}}_{l}}`$, $`\ensuremath{\mathcal{H}^{\mathrm{post}}_{l}}`$, 그리고 $`\ensuremath{\mathcal{H}^{\mathrm{res}}_{l}}`$—는 일반적인 확장 비율 $`n`$, 예를 들어 4가 입력 차원 $`C`$보다 훨씬 작기 때문에 계산 부담이 거의 없습니다. 이 설계로 인해 HC는 잔차 스트림의 정보 용량을 계층의 입력 차원에서 분리시켜, 이는 모델의 연산 복잡도(FLOPs)와 강하게 관련됩니다. 따라서 HC는 잔차 스트림 너비를 조정하는 새로운 확장 방법을 제공하며, 사전 훈련 확장 법칙에서 논의된 모델 FLOPs 및 학습 데이터 크기와 함께 사용될 수 있습니다.

그러나 HC는 잔차 스트림과 계층 입력 사이의 차원 불일치를 관리하기 위해 세 가지 매핑을 필요로 하지만, Tab. 1에 제시된 초보 실험 결과는 잔차 매핑 $`\ensuremath{\mathcal{H}^{\mathrm{res}}_{l}}`$이 가장 큰 성능 향상을 제공한다는 것을 나타냅니다. 이 발견은 잔차 스트림 내에서 효과적인 정보 교환의 중요성을 강조합니다.