BGE-M3 모델 구현하기

이 논문은 BGE-M3 모델을 Dense 레이어와 LayerNormalization만을 사용하여 TensorFlow/Keras로 처음부터 구현하는 과정을 상세히 설명합니다.

BGE-M3 소개:
BGE-M3는 70개 이상의 언어를 지원하는 다국어 임베딩 모델입니다. 약 25만 개의 토큰으로 구성된 어휘 사전을 보유하며, 특히 한국어에서 뛰어난 성능을 보입니다. MTEB 한국어 벤치마크에서 최고 수준의 성능을 달성했으며, 검색 및 분류 태스크에서 우수한 결과를 보여줍니다. 이는 기존 한국어 단일 언어 모델과도 경쟁력 있는 성능을 제공하면서 다국어 지원이라는 이점을 가집니다. 최근 RAG (Retrieval-Augmented Generation)와 같은 벡터 검색 태스크에서 활발히 활용됩니다.

3가지 Retrieval 손실 함수 구조:
BGE-M3 모델은 다음과 같은 세 가지 유형의 Retrieval 손실 함수를 동시에 최적화합니다.

Dense Retrieval: 문장 단위 CLS 벡터를 통해 의미 검색을 수행합니다. 전체 문장의 의미를 단일 벡터로 압축하여 표현합니다.

Lexical Retrieval: 토큰 단위 중요도 가중치를 학습하여 키워드 기반 검색 성능을 향상시킵니다.

Multi-Vector Retrieval: 각 토큰별 독립적인 벡터 표현을 통해 세밀한 의미 매칭이 가능하며, 토큰 단위의 벡터를 통한 검색을 지원합니다.

모델 구조 (XLMRoberta 기반):
BGE-M3의 모델 구조는 XLMRobertaModel에 기반하며, 매우 단순하고 명확합니다. 최신 Transformer 모델에서 흔히 사용되는 Rotary Position Embedding (RoPE), Pre Normalization, Linear bias 제거 등의 기법은 적용되지 않았습니다. 총 9개의 기본적인 선형 레이어 (Dense, Linear, MLP)와 3개의 LayerNormalization만으로 추론 구조를 구현할 수 있습니다. 모델은 Transformer 블록이 24번 반복되는 구조로 되어 있어, 핵심 레이어(임베딩 관련 3개, Transformer 블록 내부 9개, LayerNormalization 3개)의 구현이 중요합니다.

TensorFlow - Keras 구현:

기본 모델 클래스 정의: BGEM3TensorFlow 클래스가 tf.keras.Model을 상속받아 정의됩니다.

임베딩 레이어 구현:

* Word Embedding: tf.keras.layers.Embedding을 사용하여 250,002개의 토큰을 각각 1,024차원의 벡터로 변환합니다. $input_dim=250002$, $output_dim=1024$.
* Position Embedding: 8,194개의 위치를 각각 1,024차원의 벡터로 변환합니다. $input_dim=8194$, $output_dim=1024$.
* Token Type Embedding: BGE-M3에서는 단일 타입만 사용하므로, 1,024차원의 고정된 벡터가 모든 토큰에 동일하게 적용됩니다. $input_dim=1$, $output_dim=1024$. 추론 시 성능 최적화를 위해 미리 계산된 상수 벡터로 대체할 수 있습니다.
* LayerNormalization: $epsilon=1e-5$로 설정된 LayerNormalization 레이어가 사용됩니다.
* Forward Pass: tf.gather를 사용하여 input_ids, position_ids, token_type_ids에 해당하는 임베딩 벡터를 가져옵니다. 세 임베딩(inputs_embeds, position_embeds, token_type_embeds)을 합산한 후 layerNorm을 통과시켜 정규화합니다. position_ids는 create_position_ids_from_input_ids 함수를 통해 동적으로 생성되며, token_type_ids는 모두 0으로 채워집니다.

Transformer Block 구조 구현:

각 TransformerBlock은 tf.keras.layers.Layer를 상속받으며, 6개의 Dense 레이어, 2개의 LayerNormalization 레이어, 그리고 2번의 Residual 연산으로 구성됩니다. Dropout 레이어는 학습 시에만 사용되므로 생략 가능합니다.

* Multi-Head Attention 구현:
* Query (Q), Key (K), Value (V) 계산을 위해 self.wq, self.wk, self.wv (각각 Dense(1024)) 레이어가 사용됩니다.
* split_heads 함수를 통해 Q, K, V를 다중 헤드(num\_heads=16, depth=64)로 분리합니다.
* Scaled Dot-Product Attention: 다음 공식에 따라 attention scores를 계산합니다.
$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + \text{Mask}\right)V$$
여기서 $d_k$는 8.0으로 설정된 스케일링 인자입니다.
* Attention Mask: attention_mask는 (batch_size, 1, 1, sequence_length) 형태로 reshape되어 브로드캐스팅됩니다. 원래 마스크에서 1은 실제 토큰을, 0은 패딩을 나타내며, (1 - mask) * -10000 연산을 통해 실제 토큰 위치는 0이 되고 패딩 위치는 -10000이 되어 softmax 연산 시 해당 위치의 Attention 점수가 거의 0에 수렴하도록 합니다.
* 최종 Attention 출력은 self.dense 레이어를 통과하며, 입력(inputs)과 합산된 후 self.attlayerNorm을 통해 첫 번째 Residual 연산 및 정규화가 이루어집니다.

* Transformer Feed-Forward Neural Network (FFNN) 구현:
* intermediate Dense 레이어(Dense(4096))는 Attention 출력(attention_output)을 입력으로 받아 hidden 차원의 4배 크기로 확장합니다.
* gelu_approx 활성화 함수가 적용됩니다:
$$\text{GELU}(x) = x \cdot \frac{1}{2} \left(1 + \text{erf}\left(\frac{x}{\sqrt{2}}\right)\right)$$
* output_dense 레이어(Dense(1024))를 통해 원래 hidden 차원으로 축소됩니다.
* intermediate_output과 attention_output이 합산된 후 self.output_norm을 통해 두 번째 Residual 연산 및 정규화가 이루어집니다.

모델 Forward Flow 정리:

* 임베딩 과정에서 input_ids를 기반으로 inputs_embeds, position_embeds, token_type_embeds를 생성하고 합산 후 layerNorm을 적용하여 embedding_output을 얻습니다.
* extended_attention_mask는 (batch_size, 1, 1, sequence_length) 형태로 가공되어 Transformer 블록에 주입됩니다.
* hidden_states는 embedding_output으로 초기화된 후 24개의 encoder_layers (TransformerBlock)를 순차적으로 통과합니다.
* 출력 벡터 생성:
* Dense Retrieval: $pooled_output = hidden_states[:, 0, :]$로, 최종 hidden_states의 CLS 토큰에 해당하는 첫 번째 벡터를 사용합니다.
* Multi-Vector Retrieval: self.colbert_linear (Dense(self.d_model)) 레이어를 hidden_states[:, 1:] (CLS 토큰 제외)에 적용하고, 원본 attention_mask를 통해 패딩 토큰의 벡터를 0으로 마스킹합니다. colbert_linear의 가중치는 PyTorch 형식으로 저장된 외부 파일에서 로드하여 적용됩니다.
* 최종 출력은 dense_vecs와 colbert_vecs를 포함하는 딕셔너리 형태로 반환됩니다.

모델 Signature 생성 및 패키징 및 저장:

tf.saved_model.save를 사용하여 모델과 토크나이저를 함께 저장합니다. serving_fn을 @tf.function으로 정의하고 input_signature를 지정하여 tf.TensorSpec 형태로 입력 타입을 명시합니다. 이는 배포 가능한 SavedModel 형식을 생성하여 다양한 플랫폼(Hadoop-Spark, Spring Boot, TensorFlow-Lite, TensorFlow-Metal)에서 활용될 수 있도록 합니다.