[LLM 밑바닥부터 만들기] 32g부 - 개입: 가중치 타이잉 (Weight Tying)

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

LLM의 파라미터 효율성을 높이기 위해 입력 임베딩 층과 출력 선형 층의 가중치를 공유하는 Weight Tying 기법의 효용성을 실험했다. 입력 토큰을 벡터로 변환하는 행렬과 출력 벡터를 다시 로짓으로 변환하는 행렬을 전치 관계로 설정하여 동일한 파라미터를 재사용하는 구조이다. 실험 결과 전체 파라미터는 163M에서 124M으로 약 24% 감소했으나, 테스트 셋 손실값은 3.874로 베이스라인인 3.692보다 높게 측정되었다. 이는 가중치 공유로 인한 모델의 표현력 제한이 짧은 학습 시간 내에서는 성능 저하로 이어짐을 시사한다.

배경

Transformer 아키텍처의 기본 구조, PyTorch를 이용한 신경망 구현 경험, 임베딩 및 선형 층의 작동 원리

대상 독자

LLM 아키텍처 최적화와 파라미터 효율성에 관심 있는 AI 개발자

의미 / 영향

이 실험은 현대적인 대규모 모델들이 왜 Weight Tying을 지양하는지에 대한 실증적 근거를 제공한다. 파라미터 수 절감이라는 이득보다 모델의 표현력 제한으로 인한 성능 손실이 더 크다는 점을 명확히 보여준다.

섹션별 상세

Weight Tying은 어휘 공간과 임베딩 공간 사이의 투영이 가역적이라는 가정하에 두 층의 가중치를 공유한다. 입력 시에는 dvocab x demb 행렬을 사용하고 출력 시에는 이를 전치한 demb x dvocab 행렬을 사용하여 파라미터 효율성을 극대화한다. 이를 통해 모델 크기를 줄이고 이론적으로는 학습 속도 향상을 기대할 수 있다.

PyTorch 구현 시 nn.Embedding의 가중치를 nn.Linear의 가중치 필드에 직접 할당하는 방식으로 간단히 적용 가능하다. 하지만 두 층의 기본 초기화 방식이 정규 분포와 균등 분포로 서로 달라 가중치 공유 시 초기 손실값이 비정상적으로 높게 시작되는 현상이 발생한다. 실제 실험에서도 초기 손실이 평소 11 수준에서 460까지 치솟았으나 학습이 진행됨에 따라 점차 안정화되었다.

python

class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        // ...(중략)
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

        if cfg.get("tie_weights", False):
            self.out_head.weight = self.tok_emb.weight

PyTorch에서 입력 임베딩과 출력 헤드의 가중치를 직접 연결하여 Weight Tying을 구현하는 예시

Weight Tying 적용 시 초기 학습 손실이 비정상적으로 높게 나타나는 차트 — Chart초기화 방식이 다른 두 층의 가중치를 공유하면서 발생한 현상을 보여준다. 일반적인 초기 손실값인 11에 비해 40배 이상 높은 460 부근에서 학습이 시작됨을 확인할 수 있다.

163M 파라미터 모델에 적용한 결과 전체 파라미터의 약 1/4에 해당하는 3,800만 개 이상의 파라미터가 절감되었다. 그러나 학습 토큰 수를 동일하게 유지했음에도 불구하고 최종 학습 손실은 3.952로 베이스라인의 3.743보다 높게 나타났다. 이는 가중치 제약이 모델이 데이터를 정교하게 모델링하는 능력을 방해했음을 의미한다.

Weight Tying 모델의 전체 학습 과정 손실 곡선 — Chart초기의 급격한 하락 이후 완만한 수렴 곡선을 그리지만 최종 손실값이 베이스라인 모델보다 높게 형성됨을 시각적으로 증명한다. 최종 학습 손실은 3.952로 기록되었다.

테스트 데이터셋 평가 결과 Weight Tying 모델의 손실값은 3.874로 지금까지 테스트한 모든 개입 기법 중 가장 낮은 성능을 기록했다. GPT-2 원본 모델은 이 기법을 사용하면서도 낮은 손실을 달성했으나 이는 훨씬 더 많은 학습 에포크를 거쳐 손실을 서서히 낮췄기 때문으로 분석된다. 따라서 제한된 자원으로 빠르게 최적화해야 하는 환경에서는 부적합한 기법임이 확인됐다.

실무 Takeaway

Weight Tying을 적용하면 파라미터 수를 약 24% 절감할 수 있지만 모델의 유연성이 감소하여 동일 학습량 대비 성능이 저하될 수 있다.
서로 다른 초기화 방식을 가진 층 간에 가중치를 공유할 경우 초기 학습 불안정성이 발생하므로 초기화 전략을 일치시키는 등의 주의가 필요하다.
제한된 컴퓨팅 자원으로 단기간에 성능을 최적화해야 하는 시나리오에서는 Weight Tying보다 가중치 감쇠나 학습률 스케줄링이 더 효과적인 개입 수단이다.

언급된 리소스

문서Build a Large Language Model (from Scratch)