트랜스포머 아키텍처의 구현과 학습: 경량 생성형 언어 모델 개발 여정

들어가며: 언어 모델의 핵심을 파헤치다

현대 자연어 처리의 중심에는 트랜스포머(Transformer) 아키텍처가 있습니다. 2017년 “Attention is All You Need” 논문을 통해 소개된 이래, 트랜스포머는 NLP 분야를 완전히 재편했습니다. GPT, BERT, T5와 같은 주요 언어 모델들은 모두 이 아키텍처의 변형을 기반으로 합니다. 하지만 대규모 언어 모델의 내부 작동 방식을 진정으로 이해하기 위해서는 이론적 지식을 넘어, 직접 구현하고 학습시켜보는 경험이 필수적입니다.

본 포스팅에서는 약 6,400만 파라미터 규모의 경량 생성형 언어 모델을 처음부터 설계하고, 구현하고, 학습시킨 경험을 공유하고자 합니다. 이 과정은 단순한 코딩 실습을 넘어, 언어 모델의 내부 메커니즘과 학습 역학을 깊이 이해할 수 있는 귀중한 여정이었습니다.

모델 아키텍처: 경량화된 트랜스포머의 설계

우리의 모델은 기본적인 트랜스포머 디코더 구조를 따르되, 계산 효율성을 위해 여러 최적화를 적용했습니다. 구체적인 구조는 다음과 같습니다:

Transformer(
 (token_emb): Embedding(50257, 960)
 (pos_emb): Embedding(1024, 960)
 (blocks): ModuleList(
   (0-7): 8 x TransformerBlock(
     (attn): MultiheadAttention(
       (out_proj): NonDynamicallyQuantizableLinear(in_features=960, out_features=960, bias=True)
     )
     (norm1): LayerNorm((960,), eps=1e-05, elementwise_affine=True)
     (norm2): LayerNorm((960,), eps=1e-05, elementwise_affine=True)
     (ff1): LoRALinear(
       (dropout): Dropout(p=0.05, inplace=False)
     )
     (act): SiLU()
     (ff2): LoRALinear(
       (dropout): Dropout(p=0.05, inplace=False)
     )
     (drop): Dropout(p=0.1, inplace=False)
   )
 )
 (lm_head): Linear(in_features=960, out_features=50257, bias=True)
)

주목할 만한 설계 결정사항들은 다음과 같습니다:

1. 임베딩 차원 (960): 표준 GPT 모델들(1024 또는 그 이상)보다 약간 작은 임베딩 차원을 선택했습니다. 이는 모델 크기와 계산 요구사항을 줄이면서도 충분한 표현력을 유지하기 위한 결정이었습니다.
2. 8개의 트랜스포머 블록: 계층 깊이는 모델의 추론 능력과 직접적인 관련이 있습니다. 8개의 블록은 경량 모델에서 복잡성과 효율성 사이의 균형점을 제공합니다.
3. LoRALinear 층: 저랭크 적응(Low-Rank Adaptation) 기법을 사용하여 피드포워드 네트워크를 구현했습니다. 이는 파라미터 효율성을 높이는 동시에 미세조정 유연성을 향상시킵니다.
4. SiLU 활성화 함수: ReLU나 GELU보다 더 부드러운 특성을 가진 SiLU(Sigmoid Linear Unit, 일명 Swish)를 선택했습니다. 이는 최근 연구에서 성능 향상을 보여준 활성화 함수입니다.
5. 계층 정규화 (LayerNorm): 표준 LayerNorm을 각 서브층 이전에 적용하여 학습 안정성을 높였습니다.

이 모델 구조는 약 6,400만 개의 학습 가능한 파라미터를 가지고 있으며, 이는 현대적 기준으로는 상당히 경량화된 규모입니다. 이러한 크기는 연구 목적으로 실험하기에 이상적이며, 단일 GPU에서도 합리적인 시간 내에 학습이 가능합니다.

학습 과정: 대규모 말뭉치에서의 자기지도 학습

모델 학습은 웹 텍스트, 책, 위키백과 등 다양한 소스의 텍스트로 구성된 대규모 말뭉치를 사용했습니다. 학습 과정의 핵심 측면은 다음과 같습니다:

1. 토큰화: GPT 계열 모델과 호환성을 위해 50,257개의 어휘를 가진 BPE(Byte-Pair Encoding) 토크나이저를 사용했습니다.
2. 컨텍스트 길이: 최대 1024 토큰의 컨텍스트 윈도우를 사용하여 장거리 의존성을 포착할 수 있게 했습니다.
3. 학습 목표: 표준 자기회귀적 언어 모델링 목표, 즉 주어진 이전 토큰들을 기반으로 다음 토큰을 예측하는 방식을 사용했습니다.
4. 최적화 알고리즘: AdamW 옵티마이저를 코사인 학습률 스케줄과 함께 사용했습니다.
5. 학습 규모: 약 2백만 스텝에 걸쳐 학습을 진행했으며, 이는 대략 수십억 개의 토큰에 해당합니다.

학습 중 흥미로운 발견은 약 100만 스텝 이후에 모델이 토큰 간의 연결성과 패턴을 이해하기 시작했다는 점입니다. 이는 복잡한 언어 구조를 학습하는 데 필요한 최소한의 학습량에 대한 통찰을 제공합니다.

관찰 결과: 언어 모델의 능력 발현

약 2백만 스텝의 학습 후, 모델은 흥미로운 능력을 보여주기 시작했습니다:

1. 기본적인 언어 이해: 모델은 문법적으로 올바른 문장을 생성할 수 있게 되었습니다.
2. 문맥 인식: 주어진 프롬프트의 스타일과 주제를 유지하는 텍스트를 생성할 수 있습니다.
3. 지식 획득: 학습 데이터에 포함된 일반적인 세계 지식의 일부를 기억하고 활용하기 시작했습니다.
4. 토큰 간 연결성: 가장 주목할 만한 점은 모델이 토큰 간의 복잡한 관계를 이해하기 시작했다는 것입니다. 이는 단순히 통계적 패턴을 넘어 언어의 구조적 측면을 포착하고 있음을 시사합니다.

그러나 이 수준의 학습으로는 여전히 한계가 분명합니다. 모델은 복잡한 추론, 일관된 장문 생성, 특수 도메인 지식 등에서는 제한적인 능력을 보입니다. 이는 더 큰 모델 크기와 더 많은 학습 데이터의 필요성을 시사합니다.

향후 개선 방향: 더 나은 언어 모델을 향해

현재 모델의 성능과 학습 효율성을 더욱 향상시키기 위한 몇 가지 개선 방향을 구상하고 있습니다:

1. RMSNorm 도입: LayerNorm을 RMSNorm(Root Mean Square Layer Normalization)으로 대체하는 것을 고려 중입니다. RMSNorm은 계산 효율성이 더 높으며, 최근 연구에서 더 안정적인 학습 특성을 보여주었습니다.
2. SwiGLU 활성화 함수: 현재의 SiLU를 SwiGLU(Swish-Gated Linear Unit)로 업그레이드할 계획입니다. 이는 PaLM, Gemma와 같은 최신 모델에서 사용되는 기법으로, 더 풍부한 표현력을 제공합니다.
3. 지식 증류 적용: 현재는 Hard label(정답 토큰만 예측)을 사용한 학습 방식을 사용하고 있지만, Gemma 2b, 7b와 같이 더 큰 교사 모델의 Soft label(확률 분포)을 활용한 지식 증류(Knowledge Distillation) 학습을 시도할 계획입니다. 이 방식은 작은 모델이 더 큰 모델의 “지혜”를 효율적으로 흡수할 수 있게 합니다.
4. 임베딩 차원 최적화: 다양한 임베딩 크기 실험을 통해 모델 크기와 성능 사이의 최적 지점을 찾을 계획입니다.
5. 집중적 미세조정: 일반 사전학습 후, 특정 작업이나 도메인에 맞춘 지도 미세조정(SFT)을 적용하여 모델의 실용성을 높일 예정입니다.

학술적 의미: 모델 스케일링의 함의

이 연구는 단순히 하나의 경량 모델을 구축하는 것 이상의 의미를 가집니다. 이는 언어 모델의 능력이 어떻게 발현되는지, 그리고 어떤 요소들이 그 발전에 기여하는지에 대한 더 깊은 이해를 제공합니다.

특히 주목할 만한 점은:

1. 스케일링 법칙의 검증: 우리의 6,400만 파라미터 모델이 보여주는 능력과 한계는 “언어 모델은 규모에 따라 능력이 예측 가능하게 향상된다”는 스케일링 법칙 가설을 지지합니다.
2. 아키텍처 최적화의 중요성: 단순히 모델 크기를 키우는 것 외에도, 아키텍처 선택(활성화 함수, 정규화 기법 등)이 성능에 상당한 영향을 미친다는 점을 확인했습니다.
3. 학습 방법론의 영향: 향후 지식 증류 실험은 대규모 모델의 지식을 더 작은 모델로 효율적으로 전달하는 방법에 대한 통찰을 제공할 것입니다.

결론: 이론에서 실천으로

트랜스포머 아키텍처를 직접 구현하고 학습시키는 과정은 이론적 지식을 실제 경험으로 전환하는 귀중한 기회였습니다. 이 과정을 통해 언어 모델의 내부 작동 방식에 대한 깊은 이해를 얻을 수 있었으며, 논문에서 읽은 개념들이 실제로 어떻게 작동하는지 직접 확인할 수 있었습니다.

6,400만 파라미터의 경량 모델이지만, 약 2백만 스텝의 학습 후 보여준 능력은 놀랍습니다. 이는 언어 모델이 어떻게 지식을 습득하고, 패턴을 인식하며, 점진적으로 더 복잡한 언어적 개념을 이해하게 되는지에 대한 귀중한 통찰을 제공합니다.

향후 RMSNorm, SwiGLU, 지식 증류와 같은 최신 기법을 적용하여 모델을 더욱 발전시킬 계획입니다. 이러한 실험은 단순히 더 나은 모델을 구축하는 것을 넘어, 언어 모델의 학습 역학과 능력 발현 과정에 대한 우리의 이해를 더욱 깊게 할 것입니다.

트랜스포머 아키텍처의 구현과 학습 과정은 분명 도전적이었지만, 그 결과물과 얻은 통찰은 그만한 가치가 있었습니다. 이러한 실무적 경험은 이론적 지식을 보완하며, 더 나은 언어 모델을 설계하고 구축하는 데 필수적인 기반을 제공합니다.