The paper details the implementation of the BGE-M3 (BAAI General Embedding - Multi-lingual, Multi-task, Multi-vector) model using TensorFlow-Keras, emphasizing its lean architecture that primarily relies on Dense layers and LayerNormalization for inference. BGE-M3 is a multilingual embedding model supporting over 70 languages, noted for its strong performance in Korean MTEB benchmarks and its utility in Retrieval-Augmented Generation (RAG) tasks.

The model's core architecture is based on the XLMRobertaModel, which is characterized by its simplicity, avoiding more recent complexities like Rotary Position Embedding (RoPE), Pre-Normalization, or linear bias removal. The paper highlights that the inference structure can be realized with just nine fundamental linear layers (Dense, Linear, MLP) and three LayerNormalizations per block, repeated across 24 Transformer blocks.

The implementation consists of three main parts:

1. Embedding Layer:
   • Word Embedding: Maps 250,002 tokens to 1024-dimensional vectors. Implemented via tf.keras.layers.Embedding.
   • Position Embedding: Encodes positional information for up to 8,194 positions into 1024-dimensional vectors. Implemented via tf.keras.layers.Embedding. Position IDs are generated dynamically using create_position_ids_from_input_ids, which computes cumulative sums of a mask derived from input_ids (where padding tokens are ignored).
   • Token Type Embedding: A single 1024-dimensional constant vector applied to all tokens, as BGE-M3 uses a single token type. Implemented via $tf.keras.layers.Embedding(input_dim=1, output_dim=1024)$.
   • All three embeddings are summed: $\text{embedding\_output} = \text{inputs\_embeds} + \text{position\_embeds} + \text{token\_type\_embeds}$.
   • Finally, the combined embedding output undergoes LayerNormalization with $\epsilon = 10^{-5}$: $\text{embedding\_output} = \text{LayerNorm}(\text{embedding\_output})$.

2. Transformer Block: Each block is composed of Multi-Head Self-Attention (MHA) and a Feed-Forward Neural Network (FFNN), with residual connections and LayerNormalization applied after each sub-layer. The model uses 24 such blocks.
   • Multi-Head Attention (MHA):
     • Input inputs (from embedding or previous block) is linearly transformed into Query (Q), Key (K), and Value (V) tensors using distinct tf.keras.layers.Dense layers, each with $output_dim=1024$.
     • Q, K, V are then split into 16 heads ($num_heads=16$), where each head has a depth of $1024 / 16 = 64$. This involves reshaping and transposing: $(batch\_size, \text{seq\_len}, \text{d\_model}) \rightarrow (batch\_size, \text{num\_heads}, \text{seq\_len}, \text{depth})$.
     • Scaled Dot-Product Attention: Attention scores are computed as $\text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}} + \text{mask}\right)V$. Here, $d_k$ is the dimension of the keys (64), so the scores are divided by $\sqrt{64} = 8.0$.
     • Attention mask is applied by adding a large negative value ($-10000.0$) to padding token positions before softmax, effectively zeroing out their attention scores. The extended_attention_mask is derived from attention_mask_origin by $(1 - \text{mask}) \times -10000.0$.
     • The output from all heads is concatenated and passed through a final tf.keras.layers.Dense layer ($output_dim=1024$).
     • A residual connection adds the initial inputs to the attention output, followed by LayerNormalization: $\text{attention\_output} = \text{LayerNorm}(\text{attention\_output} + \text{inputs})$.
   • Feed-Forward Neural Network (FFNN):
     • The output from the attention sub-layer (attention_output) is passed through an intermediate tf.keras.layers.Dense layer that expands the dimension to $intermediate_size=4096$.
     • A GELU (Gaussian Error Linear Unit) approximation activation function is applied: $\text{GELU}(x) = x \times 0.5 \times (1.0 + \text{erf}(x / \sqrt{2.0}))$.
     • The result then passes through an output_dense tf.keras.layers.Dense layer, projecting back to $d_model=1024$.
     • A second residual connection adds attention_output to the FFNN output, followed by LayerNormalization: $\text{output} = \text{LayerNorm}(\text{layer\_output} + \text{attention\_output})$.

3. Model Forward Flow and Output:
   • Input input_ids and attention_mask are processed through the embedding layer and then iteratively through 24 Transformer blocks. The output of the last block is hidden_states.
   • Dense Retrieval Output: The pooled output for dense retrieval is extracted as the first token's hidden state (CLS token): $\text{pooled\_output} = \text{hidden\_states}[:, 0, :]$.
   • Multi-Vector Retrieval Output: An additional colbert_linear tf.keras.layers.Dense layer is applied to the non-CLS tokens of hidden_states (i.e., hidden_states[:, 1:]). This output is then masked by the original attention mask (excluding the CLS token) to zero out padding tokens: $\text{colbert\_vecs} = \text{colbert\_linear}(\text{hidden\_states}[:, 1:]) \times \text{attention\_mask\_origin}[:, 1:][:, :, \text{None}]$.
   • The model returns a dictionary containing dense_vecs and colbert_vecs.
   • The paper also provides instructions for loading the colbert_linear weights, which are typically found as separate PyTorch files.

The paper concludes by demonstrating how to save the implemented TensorFlow-Keras model with a serving signature for deployment, enabling its use across various platforms and applications, from large-scale Hadoop/Spark jobs to mobile inference via TensorFlow Lite.