ShellElastic이 한 번의 사전학습으로 여러 추론 예산의 모델을 만든다

무엇인가

기반 모델은 과제와 배포 플랫폼마다 다른 추론 예산을 감당해야 하는데, 예산마다 별도 모델을 학습시키면 사전학습 비용이 그대로 곱해진다. 탄성(elastic) 모델은 하나의 가중치 공유 슈퍼넷 안에 여러 중첩 서브네트워크를 함께 학습시켜 이 문제를 풀려 하지만, 기존 레시피에는 세 가지 걸림돌이 있다. 첫째, 스텝마다 서브네트워크를 샘플링하면 계산 그래프와 파라미터 접근 패턴, 작업량이 매 스텝 달라져 파라미터 샤딩 기반 대규모 분산 학습을 어렵게 만든다. 둘째, 한 크기에서 튜닝한 하이퍼파라미터가 다른 크기에서도 유효하다는 보장이 없다. 셋째, 모든 서브네트워크가 동시에 최대 업데이트(maximal update)를 받는다는 보장이 없다. 이 논문은 첫 번째를 ShellElastic 아키텍처로, 나머지 둘을 μ-Elastic 파라미터화로 해결한다.

어떻게 동작하나

ShellElastic은 히든 차원을 조밀한 코어(dense core)와 그 뒤를 잇는 셸(shell)들로 나눈다. 기저 폭 g에 대해 지원 폭 집합 S={s1<…<sK=M}을 정하면 서브네트워크 s의 폭은 sg, 전체 폭은 Mg가 되고, 이 경계들이 크기 (s_k − s_{k−1})g인 연속 그룹 G_k를 정의한다. G_1이 코어, G_2…G_K가 셸이며 서브네트워크 s_k는 앞의 k개 그룹을 포함한다. 핵심은 모든 히든 행렬을 그룹 경계 기준 블록 하삼각으로 만드는 것이다. 출력 그룹 k의 블록 행 W_k는 입력의 앞 s_k·g 차원만 읽는다(y_{G_k} = x_{:s_k g} W_k^T). 정규화는 누적 프리픽스가 아니라 그룹 내부에서 RMSNorm으로 수행하고, 어텐션은 헤드 수를 탄성적으로 늘리되 헤드 차원은 고정하며 g가 헤드 차원의 배수가 되도록 요구한다. FFN은 SwiGLU에 확장비 e를 곱한 구조로 같은 그룹 의존 패턴을 따른다. 임베딩과 리드아웃도 폭별로 따로 두어, 서브네트워크 s는 앞 sg 차원으로 로짓을 계산한다. 이 구조에서는 전체 폭 순전파 한 번이 모든 서브네트워크의 정확한 히든 상태를 프리픽스로 담는다(Proposition 1). 따라서 결합 학습의 계산 그래프가 고정되고, 스텝당 비용이 가장 넓은 서브네트워크 하나를 학습하는 것과 비슷한 수준에 머문다. 논문은 결합 학습 비용을 C_joint = C_trunk(M) + Σ C_out(s), 개별 학습 비용을 C_separate = Σ[C_trunk(s) + C_out(s)]로 비교하며, 후자는 폭마다 트렁크 패스를 한 번씩 치러야 한다고 지적한다.

무엇과 다른가

μ-Elastic은 최대 업데이트 파라미터화(μP)를 탄성 모델로 확장한 것이다. 비탄성 프록시 폭 d0에서 기저 초기화 스케일 σ0, 기저 학습률 η0, 가중치 감쇠 λ0를 얻고, 기저 폭 g에 대해 m = g/d0(μ-Transfer 인자)와 q(s) = m·s = sg/d0(μ-Internal 인자)를 정의한다. 히든 가중치의 경우 s는 출력 그룹의 지원 경계이고 그 입력 폭은 어느 서브네트워크에서든 sg로 고정되므로, 하나의 규칙이 모든 서브네트워크에 적용된다. 초기화·학습률·가중치 감쇠·순전파 승수는 활성 입력 폭에 따라 AdamW 기준으로 표로 정리되며, FFN 다운 프로젝션에는 확장된 입력 폭 esg를 반영해 √e 초기화 인자와 1/e 순전파 승수를 쓴다. 임베딩은 기저값을 그대로 쓰고, 정규화 게인과 로짓 스케일 α_s는 1로 초기화해 학습률 η0, 가중치 감쇠 0을 적용한다. 표준 파라미터화(SP)는 이 규칙에서 m=s=e=1로 두고 α_s=1로 고정한 것에 해당한다. 이론적으로 Theorem 1은 경계 s를 가진 그룹이 그것을 포함하는 모든 서브네트워크 t≥s에서 초기화 시 Θ(1) 좌표를 갖고 AdamW 한 스텝 후 Θ(η0)만큼 움직이며, 그 상수가 s, t, M, m에 독립적임을 보인다. 반대로 Proposition 2는 모든 그룹이 하나의 초기화 분산·학습률·승수를 공유하면 경계 s와 s'의 그룹이 사전활성화 크기에서 Θ(√(s'/s)), 업데이트 크기에서 Θ(s'/s)만큼 어긋나 어떤 공유값으로도 모두를 최대로 만들 수 없다고 주장한다. Corollary 1은 이 규칙이 프록시에서 튜닝한 값을 탄성 패밀리 전체와 더 넓히거나 셸을 늘린 슈퍼넷에도 그대로 통하게 한다는 두 단계 전이(μ-Transfer)를 정리한다. 같은 인자는 MatFormer, 완전 탄성 모델, StairFormer에도 적용된다.

어떻게 쓰나

실험의 주 모델은 기저 폭 g=128, S={4,5,6,7,8}인 슈퍼넷으로, 512차원 코어 뒤에 128차원 셸 네 개가 붙어 512~1024 폭의 서브네트워크 다섯 개를 만든다. 세 가지 규모(약 119M–285M, 375M–906M, 675M–1.69B 파라미터)에서 MatFormer, 완전 탄성 모델, 개별 학습 모델과 비교한다. 다섯 크기를 한 번에 학습하는 방법들 중 ShellElastic이 세 규모 모두에서 평균 검증 손실이 가장 낮았고(최선 베이스라인 대비 0.028, 0.005, 0.018 nats 낮음), 평균 5-shot 정확도도 모든 규모에서 가장 높았다. 검증 손실 15개 구간 중 12개, 5-shot 정확도 14개에서 최고이거나 공동 최고였다. 격차는 규모가 커질수록 벌어져 Low에서 최대 0.016 nats이던 것이 Max에서 0.023~0.050 nats가 되는데, 저자들은 탄성 베이스라인의 Max 전용 파라미터가 평균 다섯 스텝에 한 번만 그래디언트를 받는 반면 ShellElastic은 모든 파라미터가 매 스텝 그래디언트를 받기 때문으로 본다. 개별 학습 모델의 5분의 1 스텝만으로도 평균 검증 손실에서 0.039, 0.045, 0.014 nats 이내, 평균 5-shot 정확도에서 0.8포인트 이내에 도달한다.

전제와 한계

하이퍼파라미터 전이 실험에서 프록시의 최적값 η0=2×10⁻³, σ0=3.3×10⁻³이 세 탄성 모델 모두에 두 단계 전이로 그대로 통했다. 119M–285M 규모에서 두 값은 최적이거나 최적에 인접했고, m이나 M을 키우고 배치·스텝을 두 배로 늘린 375M–906M 규모에서도 마찬가지였으며 σ0는 MatFormer에서도 정확히 최적이 됐다. 기저값은 모든 μ-Elastic 스윕의 최적점에서 0.003 nats 이내, 모든 서브네트워크에 대해 0.005 nats 이내였다. 반면 SP는 최적점이 프록시에서 벗어나 규모에 따라 이동했고, 기저값에서 0.03~0.11 nats 더 나빴다. 가중치 감쇠도 전이되어 λ0=10⁻³에서 10⁻¹ 범위에서 평균 손실이 0.003 nats 이내로 평평했지만 SP는 모든 값에서 더 높았다. Table 3에서는 SP에 12개(375M–906M에서는 13개) 설정 중 최선을 주고 μ-Elastic에는 기저값만 줬는데도, 119M–285M에서 μ-Elastic이 모든 구간에서 더 낮은 학습·검증 손실을 보였고(평균 검증 손실 0.011, 0.013 nats), 375M–906M에서 튜닝된 SP가 세 모델 모두에서 0.006~0.021 nats 더 나빴다. 튜닝된 SP를 쓴 경우에도 ShellElastic은 두 베이스라인을 0.026~0.027, 0.020~0.024 nats 앞섰다.

StairFormer와의 비교에서 StairFormer는 g=208, S={1,2,3,4,5}로 더 좁은 서브네트워크(32.6M~200.4M)를 갖고 Max만 ShellElastic과 크기가 맞는다. 각 ShellElastic 서브네트워크의 파라미터 수에서 StairFormer의 검증 손실은 0.055~0.083 nats 높았고, 파라미터가 1.5배인 StairFormer의 Max조차 ShellElastic의 High보다 0.001 nats 높았다. Max에서 평균 5-shot 정확도는 36.9 대 35.7이었다. 아키텍처 탐색에서는 히든 행렬이 M² 블록 중 B(S)=Σ_k (s_k−s_{k−1})s_k개를 유지한다는 점을 이용해, 파라미터 수와 폭 개수를 고정하고 넓은 그룹을 어디에 둘지 비교했다. B=42, 폭 다섯 개에서는 폭 g 그룹 넷과 4g 그룹 하나가 되는데, 넓은 그룹을 코어에 두는 배치가 모든 공유 폭에서 최저 손실이거나 0.006 nats 이내였고 셸에 두면 그 넓은 그룹을 포함하는 첫 폭이 0.025~0.052 nats 나빠졌다. 이것이 넓은 코어와 얇은 셸이라는 설계를 뒷받침한다.

절제 실험에서 서브네트워크 선택, 인플레이스 증류, 서브네트워크 앙상블, 스테이지-잔차 그래디언트 스케일링, 좁은 폭의 트렁크 그래디언트는 평균 검증 손실을 0.010~0.20 nats 바꿨고, 손실 가중치·정규화·리드아웃·기저 폭과 헤드 차원·최적화 하이퍼파라미터는 0.007 nats 이내만 바꿨다. 매 스텝 다섯 개를 모두 학습하는 All이 기본값으로 평균 손실이 다른 규칙보다 0.012~0.059 nats 낮았고, Random은 Max에서 0.095 nats 나빴다. 인플레이스 증류는 Low를 최대 0.034 nats 개선하고 평균을 0.010 nats 낮춰 절제 중 가장 큰 이득을 줬지만 Max를 0.008 nats 악화시키고 스텝 시간을 2.2~2.6배로 늘려 채택되지 않았다. 앙상블은 어떤 방식도 Max 단독을 이기지 못했고, 라우터는 50스텝 안에 모든 가중치를 Max에 몰아 다른 서브네트워크가 최대 0.365 nats 손실을 봤다. 다만 정답을 아는 오라클은 Max보다 0.19 nats 낮았고 토큰의 71%에서 더 좁은 서브네트워크가 더 나았는데, 저자들은 더 나은 앙상블이 거의 추가 비용 없이 ShellElastic을 개선할 유망한 방향이라고 본다.

한계와 전제도 분명하다. μ-Elastic은 초기화와 업데이트의 스케일을 고정할 뿐 최적 하이퍼파라미터 값을 보장하지 않으며, 폭이 유일한 탄성 축이다. 깊이, 미니배치, 학습 스텝 수는 모든 서브네트워크가 공유하므로 기존의 깊이·배치·토큰 예산 스케일링 규칙을 따른다. 이론적 보장은 μP의 표준적 휴리스틱 가정 아래에서 증명된다. 또한 탄성 모델은 다른 언어 모델과 마찬가지로 학습 데이터의 편향을 물려받는다. 실무 관점에서 이 논문이 주는 실질적 지침은 명확하다. 작은 비탄성 프록시에서 학습률·초기화·가중치 감쇠만 한 번 튜닝하고, 기저 폭과 지원 폭 집합을 정한 뒤 μ-Elastic 규칙으로 스케일을 계산하면, 매 스텝 모든 폭을 함께 학습시키면서도 스텝당 비용은 가장 넓은 모델 하나와 비슷하게 유지할 수 있다. 다만 넓은 그룹을 코어에 두는 배치, 그룹 내부 정규화, 폭별 리드아웃, 그리고 인플레이스 증류를 쓰지 않는 선택이 모두 실측 근거에 기반한다는 점을 확인하고 자기 워크로드에 맞게 검증해야 한다.