[Paper] 가중치 공유가 토큰 임베딩을 출력 공간으로 편향시킴

Source: arXiv - 2603.26663v1

개요

Weight tying—입력 토큰 임베딩과 출력 “언임베딩” 투영에 동일한 행렬을 사용하는 방법—은 현대 언어 모델 아키텍처의 핵심 요소입니다. 이 논문은 숨겨진 편향을 밝혀냅니다: 공유 행렬이 주로 출력 측면에 최적화되어 있어, 하위 레이어가 의존하는 입력 표현의 품질을 저하시킬 수 있습니다. 이러한 편향을 이해하면, 특히 임베딩 행렬이 전체 파라미터의 큰 비중을 차지하는 작은 모델에서, 가중치 묶기가 성능에 악영향을 미치는 이유를 설명할 수 있습니다.

주요 기여

• 경험적 증거: 연결된 임베딩 행렬이 동일한 규모의 비연결 모델의 입력 임베딩보다 출력(언임베딩) 행렬과 더 가깝게 정렬된다는 사실.
• 그라디언트 분석: 초기 학습 단계에서 출력 그라디언트가 지배적이며, 공유 행렬을 출력 작업 쪽으로 끌어당긴다는 점을 보여줌.
• 튜닝 렌즈 진단: 가중치 연결을 사용할 경우 초기 트랜스포머 층이 더 약하고 덜 유용한 신호를 받는다는 것을 밝혀냄.
• 중재 실험: 학습 중 입력 측 그라디언트를 확대하면 편향이 감소하고 하위 층의 기여도가 향상됨.
• 실용적 통찰: 모델 규모가 커질수록, 특히 파라미터 효율적인 LLM에서 가중치 연결이 성능 병목이 될 수 있는 이유를 제시.

방법론

1. Model Setup – 저자들은 두 종류의 트랜스포머 언어 모델을 학습한다: 하나는 표준 가중치 공유(공유 임베딩/언임베딩 행렬)를 사용하고, 다른 하나는 입력과 출력 행렬을 별도로 두는(비공유) 모델이다. 다른 모든 하이퍼파라미터는 동일하게 유지한다.
2. Alignment Metrics – 공유 행렬과 비공유 기준선의 해당 입력/출력 행렬 사이의 코사인 유사도와 정준 상관(canonical correlation)을 계산하여 임베딩이 각 측에 얼마나 “가까운지”를 정량화한다.
3. Gradient Flow Analysis – 손실로부터 임베딩 행렬로 흐르는 그래디언트 크기를 추적함으로써, 초기 학습 단계에서 출력 관련 그래디언트가 10배 정도 더 크게 나타남을 보여준다.
4. Tuned Lens Probing – 가벼운 선형 프로브(“튜닝 렌즈”)를 각 트랜스포머 레이어에 부착하여 각 레이어가 최종 예측에 얼마나 유용한 정보를 제공하는지 측정한다.
5. Gradient Scaling Intervention – 학습 중에 입력 측 그래디언트 성분에 1보다 큰 계수를 곱해 그래디언트 기여도를 재조정하고, 정렬 및 레이어 유용성에 미치는 하위 효과를 관찰한다.

Results & Findings

• Alignment Shift: 몇 백 번의 학습 단계 후, tied matrix의 방향이 untied 모델의 input matrix보다 output matrix에 약 30 % 더 가깝게 정렬됩니다 (코사인 유사도 기준).
• Gradient Imbalance: 학습 초기 10 % 구간에서 output gradient가 input gradient보다 약 5–10배 크게 나타나, 공유된 매트릭스를 출력 공간 쪽으로 끌어당깁니다.
• Layer Utility Drop: tied 모델의 초기 transformer 레이어들은 untied 모델에 비해 residual stream에 약 15 % 적게 기여합니다 (tuned‑lens 정확도 기준). 이는 중간 표현이 약해졌음을 의미합니다.
• Bias Mitigation: input gradient를 3배로 스케일링하면 정렬 균형이 회복되어 (input vs. output 유사도가 대략 동일해짐) 초기 레이어 유틸리티 손실의 약 8 %가 회복됩니다. 이는 퍼플렉시티가 약 0.3–0.5 % 낮아지는 일관된 개선 효과로 이어집니다.

실용적 함의

• Model Scaling Decisions – 전체 파라미터 중 임베딩 행렬이 아주 작은 부분을 차지하는 대형 LLM의 경우, 이 편향은 무시될 수 있습니다. 그러나 소‑중형 모델(예: 1 억–10 억 파라미터)에서는 임베딩 행렬이 메모리와 연산을 지배할 수 있어, 파라미터를 공유하면 의도치 않게 표현 품질이 떨어질 수 있습니다.
• Fine‑Tuning Strategies – 다운스트림 작업에 대해 파라미터가 공유된 모델을 미세 조정할 때는 임베딩을 풀링(untying) 하거나 그래디언트 균형 스케줄을 적용해 입력 측이 따라잡을 수 있도록 고려하십시오.
• Architecture Design – 설계자는 부분 공유(partial tying)(예: 일부 차원만 공유) 혹은 **듀얼‑헤드 임베딩(dual‑head embeddings)**을 도입하여 입력과 출력 서브스페이스를 별도로 유지하면서도 파라미터를 절감할 수 있습니다.
• Training Optimizations – 그래디언트 스케일링, 임베딩 레이어에 집중된 러닝레이트 워밍업, 혹은 보조 재구성 손실(auxiliary reconstruction losses) 같은 간단한 트릭을 사용하면 파라미터 수를 늘리지 않고도 편향을 완화할 수 있습니다.
• Interpretability Tools – 여기서 보여준 tuned‑lens 방법론은 다른 구조적 단축(예: low‑rank adapters)이 은닉 표현에 미치는 부작용을 감사(audit)하는 데 재활용될 수 있습니다.

제한 사항 및 향후 연구

• 본 연구는 표준 언어 모델링 목표로 학습된 decoder‑only transformers에 초점을 맞추고 있으며; encoder‑only 또는 encoder‑decoder 아키텍처에서는 결과가 다를 수 있습니다.
• 실험은 English corpora와 비교적 제한된 모델 크기 범위에만 한정되었으며; 2 B 파라미터를 초과하는 스케일링 추세는 아직 추측에 불과합니다.
• 그라디언트‑스케일링 개입은 hand‑crafted heuristic이며; 최적의 균형 스케줄을 자동으로 학습하는 것이 더 견고할 수 있습니다.
• 향후 연구에서는 dynamic tying(학습 중에 연결/비연결 전환)이나 입력‑출력 정렬 대칭을 명시적으로 강제하는 regularization terms를 탐색할 수 있습니다.

Bottom line: Weight tying은 공짜가 아닙니다. 파라미터를 절감하는 동시에 공유 행렬을 입력 표현보다 출력 예측에 치우치게 합니다—특히 작은 모델에서 그렇습니다. 이 그라디언트 불균형을 인식하고 수정함으로써 개발자는 티잉의 효율성 이점을 유지하면서 다운스트림 작업을 위한 풍부한 토큰 임베딩을 보존할 수 있습니다.

저자

• Antonio Lopardo
• Avyukth Harish
• Catherine Arnett
• Akshat Gupta

논문 정보

• arXiv ID: 2603.26663v1
• 카테고리: cs.CL
• 출판일: 2026년 3월 27일
• PDF: PDF 다운로드