Introduction

이 글에서는 BERT의 임베딩 레이어 구현 세부 사항에 대해 설명하겠습니다. BERT의 임베딩 레이어는 주로 토큰 임베딩(Token Embeddings), 세그먼트 임베딩(Segment Embeddings), 그리고 위치 임베딩(Position Embeddings)으로 구성되어 있습니다.

Overview

이 글에서는 BERT의 임베딩 레이어 구현 세부 사항에 대해 설명할 것입니다. 구체적으로는 토큰 임베딩(Token Embeddings), 세그먼트 임베딩(Segment Embeddings), 그리고 위치 임베딩(Position Embeddings)에 대해 다룰 것입니다.

다음은 각 임베딩 레이어의 기능을 적절히 설명한 논문의 다이어그램입니다:

대부분의 NLP 관련 작업을 해결하기 위한 딥러닝 모델처럼, BERT도 각 입력 토큰(입력 텍스트의 단어)을 토큰 임베딩(Token Embedding) 레이어를 통해 전달하여 각 토큰이 벡터 표현으로 변환됩니다. 다른 딥러닝 모델과 달리, BERT는 세그먼트 임베딩(Segment Embeddings)과 위치 임베딩(Position Embeddings)의 형태로 추가적인 임베딩 레이어를 가지고 있습니다. 이러한 추가 임베딩 레이어의 이유는 이 글의 끝부분에서 명확해질 것입니다.

Token Embeddings

Purpose

앞서 언급한 것처럼, 토큰 임베딩 레이어의 역할은 단어를 고정된 차원의 벡터 표현으로 변환하는 것입니다. BERT의 경우, 각 단어는 768차원 벡터로 표현됩니다.

Implementation

입력 텍스트가 “I like strawberries”라고 가정해봅시다. 이 입력 텍스트를 BERT 로 처리할 때, 다음과 같은 단계가 진행됩니다:

1. 토큰화:
먼저 입력 텍스트는 개별 토큰으로 토큰화됩니다. BERT는 WordPiece 토큰화를 사용합니다. 간단히 “I like strawberries”가 다음과 같이 토큰화된다고 가정해봅시다:

["[CLS]", "I", "like", "strawberries", "[SEP]"]

2. 토큰 임베딩(Token Embeddings):
각 토큰은 토큰 임베딩 레이어를 통해 768차원의 벡터 표현으로 변환됩니다. 이 벡터들을 다음과 같이 나타낼 수 있습니다:

[T_CLS, T_I, T_like, T_strawberries, T_SEP]

3. 세그먼트 임베딩(Segment Embeddings):
BERT 는 문장 쌍을 처리할 수 있으며, 세그먼트 임베딩은 서로 다른 문장에서 온 토큰들을 구분하는 데 도움이 됩니다. 이 경우 하나의 문장만 있으므로 모든 토큰에 동일한 세그먼트 임베딩이 사용됩니다. 예를 들어, S_0 세그먼트 임베딩 벡터가 모든 토큰에 사용된다고 가정합니다:

[S_0, S_0, S_0, S_0, S_0]

4. 포지션 임베딩(Position Embeddings):
토큰의 순서를 캡처하기 위해 포지션 임베딩이 추가됩니다. 각 토큰은 문장에서의 위치에 따라 포지션 임베딩이 할당됩니다. 예를 들어, 포지션 임베딩이 다음과 같다고 가정합니다:

P_1, P_2, P_3, P_4]

5. 합산:
각 토큰의 최종 임베딩은 해당 토큰 임베딩, 세그먼트 임베딩, 그리고 포지션 임베딩을 합산하여 얻습니다. 따라서 “I like strawberries” 문장의 토큰 임베딩은 다음과 같이 됩니다:
E_CLS = T_CLS + S_0 + P_0 E_I = T_I + S_0 + P_1 E_like = T_like + S_0 + P_2 E_strawberries = T_strawberries + S_0 + P_3 E_SEP = T_SEP + S_0 + P_4

이 임베딩들은 BERT 모델을 통해 전달되어 각 토큰의 문맥적 표현을 생성하게 되며, 이는 분류, 번역 또는 질문-응답과 같은 다양한 자연어 처리 작업에 사용될 수 있습니다.

예시 그림을 보겠습니다.

정리해보자면, 토큰 임베딩(Token Embeddings) 레이어는 단어를 고정된 차원의 벡터 표현으로 변환합니다. BERT의 경우, 각 단어는 768차원의 벡터로 표현됩니다.
세그먼트 임베딩(Segment Embeddings) 레이어는 입력 텍스트가 두 개의 문장으로 구성될 때, 각 문장을 구분하기 위해 사용됩니다. 각 토큰에 대해 세그먼트 임베딩 값을 추가하여 모델이 서로 다른 문장임을 인식할 수 있도록 합니다.포지션 임베딩(Position Embeddings) 레이어는 문장 내에서 각 단어의 위치 정보를 제공합니다. 이는 Transformer 모델의 특성상 순서 정보가 없기 때문에, 각 단어의 위치 정보를 임베딩 벡터에 추가하여 순서를 인식할 수 있도록 합니다.

이제 위 예시에서 보았던 “I like strawberries”라는 입력 텍스트를 예로 들어, BERT의 임베딩 레이어가 어떻게 구현되는지 코드로 작성해보겠습니다.

import torch
import torch.nn as nn

class BERTEmbeddings(nn.Module):
    def __init__(self, vocab_size, hidden_size, max_position_embeddings, type_vocab_size):
        super(BERTEmbeddings, self).__init__()
        self.token_embeddings = nn.Embedding(vocab_size, hidden_size)
        self.position_embeddings = nn.Embedding(max_position_embeddings, hidden_size)
        self.segment_embeddings = nn.Embedding(type_vocab_size, hidden_size)
        
        self.LayerNorm = nn.LayerNorm(hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(0.1)
        
    def forward(self, input_ids, token_type_ids=None):
        # Token Embeddings
        token_embeddings = self.token_embeddings(input_ids)
        
        # Position Embeddings
        seq_length = input_ids.size(1)
        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
        position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        
        # Segment Embeddings
        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_ids)
        segment_embeddings = self.segment_embeddings(token_type_ids)
        
        # Summing Embeddings
        embeddings = token_embeddings + position_embeddings + segment_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        
        return embeddings

# 예제 입력
input_ids = torch.tensor([[101, 1045, 2066, 13121, 102]])  # "I like strawberries"의 토큰 ID
token_type_ids = torch.tensor([[0, 0, 0, 0, 0]])  # 모든 토큰이 동일한 세그먼트에 속함

# BERT 임베딩 레이어 인스턴스 생성
bert_embeddings = BERTEmbeddings(vocab_size=30522, hidden_size=768, max_position_embeddings=512, type_vocab_size=2)

# 임베딩 레이어를 통해 입력을 변환
embeddings = bert_embeddings(input_ids, token_type_ids)
print(embeddings)

이 코드에서는 BERT의 세 가지 임베딩 레이어를 구현하고, “I like strawberries”라는 문장의 토큰을 벡터로 변환합니다. BERTEmbeddings 클래스는 토큰 임베딩, 포지션 임베딩, 세그먼트 임베딩을 포함하며, 입력 텍스트의 토큰 ID와 세그먼트 ID를 받아 임베딩 벡터를 반환합니다.

그림으로 나타내보면 다음과 같습니다.

이후 통과된 Embedding layer는 Encoder내에서 다음과 같이 작용합니다.

마치며

위의 과정을 통해 BERT는 입력 텍스트를 효과적으로 벡터로 변환하여, 다양한 자연어 처리 작업에 활용할 수 있는 문맥적 표현을 생성합니다. 토큰 임베딩, 세그먼트 임베딩, 포지션 임베딩의 조합은 BERT가 텍스트의 구조와 의미를 이해하는 데 중요한 역할을 하며, 이를 통해 BERT는 높은 정확도의 자연어 처리 성능을 제공합니다.

Reference

[NLP] How does BERT work?

BERT is a pre-trained model released by Google in 2018, and has been used a lot so far, showing the highest performance in many NLP tasks. In this post, let’…

Why BERT has 3 Embedding Layers and Their Implementation Details

In this article, I will explain the implementation details of the embedding layers in BERT, namely the Token Embeddings, Segment Embeddings, and the Position Embeddings. Here’s a diagram from the…

[Pytorch][BERT] 버트 소스코드 이해_⑤ BertEmbedding

Pytorch][BERT] 버트 소스코드 이해 목차 BERT 📑 BERT Config 📑 BERT Tokenizer 📑 BERT Model 📑 BERT Input 📑 BERT Output 📑 BERT Embedding 👀 📑 BERT Pooler 📑 BERT Enocder 📑 BERT Layer 📑 BERT SelfAttention 📑 BERT SelfOtput 1. BertEmbedding 이란?