기술 요약: 예측 코딩 네트워크의 무한 너비 및 깊이 한계에 대하여

문제 정의

예측 코딩(Predictive Coding, PC)은 뉴런이 가중치를 업데이트하기 전에 국소적 예측 오차(에너지 함수)를 최소화하는 생물학적으로 타당한 역전파(Backpropagation, BP)의 대안이다. 최근의 재매개변수화(reparameterization)는 심층 PC 네트워크(PCN)의 훈련 안정성을 개선했지만, 이러한 방법들을 거대한 너비와 깊이로 확장하기 위한 이론적 근거는 여전히 불분명하다. 구체적으로, PC가 무한 너비 및 무한 깊이 한계에서 BP의 안정성과 특징 학습 능력을 따라잡을 수 있는지, 아니면 표준적인 PCN이 모델 크기가 증가함에 따라 불안정성(예: 출력 폭발)을 겪게 될지는 알려져 있지 않다. 또한, 이러한 한계 조건에서 PC 그래디언트가 BP 그래디언트로 수렴하기 위한 조건에 대한 엄밀한 규명이 필요하다.

방법론

저자들은 무한 너비($N \to \infty$)와 무한 깊이($L \to \infty$)의 결합 한계에서 PCN에 대한 이론적 및 경험적 분석을 수행한다.

1. 이론적 프레임워크:

   • 일반적 매개변수화: 본 연구는 선형 다층 퍼셉트론(MLP) 및 잔차 네트워크(Residual Networks)에 대해 사전 활성화($a_\ell$), 가중치 초기화 분산($b_\ell$), 학습률($c$), 네트워크 출력($d$)의 스케일링 지수를 정의하는 일반적인 "abcd" 매개변수화 체계를 채택한다.
   • 평형 에너지(Equilibrated Energy): 선형 네트워크의 경우, 저자들은 활동(activity)에 대해 PC 에너지를 최소화한 값인 "평형 에너지"($F^*$)를 활용한다. 이들은 선형 네트워크에서 $F^*$가 표준 평균 제곱 오차(MSE) 손실의 재스케일링된 버전, 즉 $F^*(\theta) = L(\theta) / s(\theta)$임을 보여준다 (여기서 $s(\theta)$는 가중치 의존적 재스케일링 인자이다).
   • 설계 요구사항(Desiderata): 분석은 다음 세 가지 주요 기준(BP 문헌에서 차용한 것)을 만족하는 매개변수화를 찾고자 한다:
     1. 안정성(Stability): 층별 사전 활성화와 네트워크 예측이 $N \to \infty$일 때 유한하게 유지되어야 한다(사라지거나 폭발하지 않아야 함).
     2. 특징 학습(Feature Learning): 특징이 훈련 중에 비자명하게(non-trivial, non-lazy learning) 진화해야 한다.
     3. 깊이 안정성(ResNet의 경우): 은닉 활성화가 $L \to \infty$일 때 안정적으로 유지되어야 하며, 각 잔차 블록은 특징 진화에 $\Theta(1/L)$만큼 기여해야 한다.

2. 유도(Derivation):

   • 저자들은 위 설계 요구사항을 충족하기 위해 필요한 구체적인 스케일링 지수를 유도한다.
   • 재스케일링 인자 $s(\theta)$와 평형 에너지의 그래디언트의 점근적 거동을 분석한다.
   • 이러한 안정적인 매개변수화 하에서, PC 그래디언트가 너비가 커짐에 따라 BP 그래디언트로 수렴함을 증명하며, 이때 오차 항은 깊이 대 너비 비율($L/N$)에 의존한다.

3. 경험적 검증:

   • 선형 및 비선형 아키텍처(MLP, CNN, Transformer 포함)에 대해 실험을 수행하였다.
   • 토이 태스크부터 CIFAR-10, Fashion-MNIST, ImageNet-1k에 이르는 데이터셋으로 모델을 훈련하였다.
   • 본 연구는 다양한 너비와 깊이에 따른 PC 그래디언트(활동 평형 상태에서 계산됨)와 BP 그래디언트 사이의 코사인 유사도를 측정하였다.

주요 기여

1. 안정적 매개변수화의 동등성:
   본 논문은 PC를 위한 안정성과 비자명한 특징 학습을 모두 보장하는 매개변수화 집합이 BP의 그것과 동일함을 증명한다. 구체적으로, 선형 MLP 및 잔차 네트워크에 대해, "평균장(Mean-Field)" 또는 "최대 업데이트($\mu$P)" 매개변수화가 PC의 안정성 및 특징 학습 요구사항을 만족하는 유일한 해이다. 이는 표준 PCN(PyTorch의 기본값과 같은 표준 매개변수 사용)이 거대한 너비에서 본질적으로 불안정하며 출력이 폭발하게 된다는 것을 의미한다.

2. 역전파로의 수렴:
   저자들은 이러한 안정적인 매개변수화 하에서, 네트워크 너비가 깊이보다 훨씬 커질 때($L/N \to 0$) PC에 의해 계산된 그래디언트가 BP의 그래디언트로 수렴함을 입증한다.

   • 고정된 깊이의 선형 MLP의 경우, PC와 BP 그래디언트 간의 상대 오차는 $O(N^{-1})$로 스케일링된다.
   • 결합 한계에서의 선형 잔차 네트워크의 경우, 상대 오차는 $O(L/N)$으로 스케일링된다.
   • 이 수렴은 목적 함수와 그래디언트 모두에 적용되며, 이는 넓은 한계에서 PC가 효과적으로 BP를 구현함을 의미한다.

3. "Fast Saddle-to-Saddle" 영역의 불안정성:
   본 논문은 (깊고 좁은 네트워크에서) PC가 BP보다 더 빠른 "saddle-to-saddle" 역학을 보이는 것으로 관찰되었던 이전의 영역이 표준 매개변수화(SP)에 의존하고 있음을 식별한다. 저자들은 이 영역이 너비와 깊이 모두에 대해 불안정하며, 안정성 요구사항을 위반한다는 것을 보여준다. 따라서 이러한 속도 향상은 대규모 모델에서의 수치적 불안정성을 대가로 얻은 것이다.

4. 경험적 일반화:
   선형 네트워크를 바탕으로 도출된 이론적 분석에도 불구하고, 실험 결과 비선형 아키텍처(CNN, Transformer)에서도 높은 그래디언트 정렬이 유지됨을 보여주었다. 이는 활동 평형에 도달하고 적절한 최적화 도구(Adam)와 손실 함수를 사용할 때 성립한다.

결과

• 이론적: 유도 과정은 $d=1/2$라는 출력 스케일링 지수가 BP에서와 마찬가지로 PC의 non-lazy 특징 학습에 결정적임을 확인한다. 재스케일링 인자 $s(\theta)$는 $N \to \infty$ (MLP의 경우) 또는 $L/N \to 0$ (ResNet의 경우)일 때 1에 수렴하며, 이로 인해 PC 목적 함수가 MSE 손실로 수렴하게 된다.
• 경험적:
  • 그래디언트 정렬: 특히 $L/N$이 작은 ResNet 및 Transformer의 경우, 너비가 증가함에 따라 PC와 BP 그래디언트 사이의 코사인 유사도가 1.0에 근접한다.
  • 비선형 모델: 32계층 잔차 MLP, 10계층 CNN, 12블록 Transformer에서 높은 정렬이 관찰되었다.
  • 표준 매개변수화의 실패: 표준 매개변수(SP)를 사용하는 모델은 너비가 증가함에 따라 출력이 폭발하고 그래디언트 정렬이 저하되는 모습을 보였으며, 이는 이론적 불안정성을 확인시켜 주었다.

의의 및 주장

본 논문은 PC가 확장 가능한 매개변수화에 대한 최초의 원리적 제약을 제공한다고 주장한다. 주요 의의는 두 가지 측면에서 나타난다:

1. 확장성 제약: PC가 대규모 모델로 확장되기 위해서는 BP와 동일한 매개-변수화 규칙(구체적으로 $\mu$P)을 채택해야 함을 확립한다. 이는 최근의 심층 PC 연구에서 사용된 "휴리스틱"한 재매개변수화들이 생존하기 위해서는 이러한 이론적 한계와 일치해야 함을 시사한다.
2. BP의 생물학적 타당성: 이 결과는 뇌(아마도 깊이보다 훨씬 넓은 네트워크인)가 국소적 학습 규칙만을 사용하여 BP의 전역적 신용 할당(global credit assignment)을 어떻게 구현할 수 있는지에 대한 메커니즘을 제시한다. 만약 뇌가 $L/N \to 0$인 영역에서 작동한다면, PC는 이론적으로 BP로 수렴할 수 있으며, 이를 통해 비국소적 정보 전파 없이도 효율적인 학습을 가능하게 할 수 있다.

저자들은 이론적 증명이 선형 네트워크와 정확한 활동 수렴에 의존한다는 점을 언급하며 한계점에 대해 겸허한 태도를 유지한다. 또한 비선형 네트워크에서 활동 평 equilibrium에 도달하는 데 드는 계산 비용이 여전히 실질적인 병목 현상임을 인정하면서도, 향후 아날로그 하드웨어 구현이 이를 해결할 수 있을 것이라고 제안한다. 저자들은 PC가 모든 영역에서 BP보다 우월하다고 주장하는 것이 아니라, 올바른 매개변수화가 적용된 넓은 한계에서 두 알고리즘이 수학적으로 동등해진다는 점을 강조한다.