DeepLoop: 루프형 Transformer의 깊이 확장을 안정화하는 방법

도입

Looped Transformers는 저장해야 하는 파라미터 수를 늘리지 않고 순차 계산 깊이를 키우려는 접근입니다. 많은 Transformer 층을 각각 따로 두는 대신, 작은 물리 블록 묶음을 여러 라운드 반복 적용합니다. 이렇게 하면 실제 저장 파라미터는 그대로이면서, 계산상으로 펼쳐진 깊이(unrolled depth)는 더 깊어집니다.

DeepLoop 논문은 이 구조에서 자주 간과되는 안정성 문제를 다룹니다. 루프형 Transformer는 단순히 일반 Transformer의 층 수를 늘린 버전이 아닙니다. 같은 블록이 여러 번 방문되면, 공유 파라미터는 반복 방문에서 그래디언트를 모아 업데이트되고, 다음 선형화된 순전파에서 다시 같은 방문들에 의해 읽힙니다. 이 결합 때문에 잔차 스케일링 규칙도 달라져야 합니다.

핵심 내용

• 파라미터 재사용이 핵심 변수: untied Transformer에서는 각 잔차 분기가 자기만의 파라미터 업데이트를 갖습니다. 반면 looped Transformer에서는 하나의 공유 업데이트가 여러 방문에 의해 함께 사용됩니다.
• 묶인 깊이 효과의 정식화: 저자들은 1차 섭동 경계를 통해 이 현상을 설명하고, 방문 간 정렬 정도를 나타내는 visit-alignment coefficient κ_R을 도입합니다.
• DeepNorm 규칙의 한계: 방문들이 서로 탈상관되는 경우에는 기존 DeepNorm 지수를 회복할 수 있습니다. 그러나 보수적인 정렬 상황에서는 물리 깊이를 고정한 채 루프 수가 늘어날수록 지수가 1/4에서 1/2로 커져야 한다고 분석합니다.
• DeepLoop의 설정: 새로운 블록을 추가하기보다 Post-LN DeepNorm 아키텍처를 유지하고, 펼쳐진 깊이 N에 대해 α=(2N)^{1/2}, β=(8N)^{-1/2}를 사용합니다.
• 실험 결과: GPT-2 small 및 GPT-2 medium 규모의 GPT 스타일 루프형 언어 모델에서 평가했습니다. 물리 블록이 재방문되지 않을 때는 거의 중립적이지만, recurrent depth가 활성화되면 검증 손실과 다운스트림 정확도가 개선됩니다.

의미와 영향

DeepLoop의 핵심 메시지는 루프형 모델의 깊이를 단순히 층 수로만 판단해서는 안 된다는 것입니다. 같은 파라미터가 몇 번 방문되는지, 그리고 방문 간 업데이트 방향이 얼마나 정렬되는지가 안정성에 큰 영향을 줍니다.

이는 파라미터 효율적인 LLM, 반복 계산 구조, 테스트 타임 계산 확장과 관련된 연구에 중요합니다. 작은 블록을 반복 사용해 능력을 끌어내려면, 아키텍처뿐 아니라 잔차 스케일링과 초기화 규칙도 반복 구조에 맞춰 조정되어야 합니다.

물론 현재 결과는 GPT 스타일 모델의 GPT-2 small 및 medium 규모에 기반합니다. 더 큰 모델, 더 긴 루프, 다른 학습 레시피에서도 같은 효과가 유지되는지는 추가 검증이 필요합니다. 그럼에도 DeepLoop는 안정적인 순환 깊이를 위해서는 명목상 층 수가 아니라 파라미터 방문 패턴을 고려해야 한다는 분명한 방향을 제시합니다.

출처: Hugging Face Daily Papers