[LLM 밑바닥부터 만들기] 32g부 - 개입: 가중치 타이잉 (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 원본 모델은 이 기법을 사용하면서도 낮은 손실을 달성했으나 이는 훨씬 더 많은 학습 에포크를 거쳐 손실을 서서히 낮췄기 때문으로 분석된다. 따라서 제한된 자원으로 빠르게 최적화해야 하는 환경에서는 부적합한 기법임이 확인됐다.

용어 해설

가중치 타이잉(Weight Tying): — 입력 임베딩 층과 출력 선형 층의 가중치 행렬을 공유하는 기법이다. 모델의 파라미터 수를 획기적으로 줄일 수 있으나, 입력과 출력의 임베딩 공간이 동일해야 한다는 제약이 생겨 모델의 유연성이 제한될 수 있다.
임베딩 행렬(Embedding Matrix): — 이산적인 토큰 ID를 고차원의 연속적인 벡터로 변환하는 가중치 집합이다. 각 행은 특정 토큰의 의미적 표현을 담고 있으며, 모델 학습의 출발점이 되는 중요한 파라미터이다.
가중치 초기화(Initialization): — 학습 시작 전 모델의 파라미터에 초기값을 부여하는 과정이다. 층의 종류에 따라 정규 분포나 균등 분포 등 서로 다른 통계적 방식을 사용하며, 부적절한 초기화는 학습 초기 단계의 손실값을 비정상적으로 높일 수 있다.
손실 차트(Loss Chart): — 학습 진행에 따른 모델의 오차 변화를 시각화한 그래프이다. 모델이 데이터의 패턴을 얼마나 잘 학습하고 있는지 보여주며, 수렴 속도와 최종 성능을 판단하는 핵심 지표로 활용된다.

기술

PyTorch
Python
GPT-2
nn.Embedding
nn.Linear

활용 사례

경량 LLM 설계
임베딩 공간 최적화
모델 파라미터 절감 실험

언급된 리소스

문서Build a Large Language Model (from Scratch)