transformer_decoder.py

import torch
import torch.nn as nn
from transformers import PretrainedConfig


class CustomDecoderConfig(PretrainedConfig):
    model_type = "custom-transformer-decoder"

    def __init__(
        self,
        vocab_size=30522,
        max_position_embeddings=128,
        hidden_size=256,
        num_attention_heads=8,
        intermediate_size=1024,
        num_hidden_layers=4,
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        **kwargs
    ):
        super().__init__(**kwargs)

        # 어휘 크기(vocab size): 입력 token id의 전체 개수
        self.vocab_size = vocab_size

        # 최대 위치 임베딩(max position embeddings): 입력 sequence의 최대 길이
        self.max_position_embeddings = max_position_embeddings

        # 은닉 차원(hidden size): 각 token vector의 차원
        self.hidden_size = hidden_size

        # attention head 개수(num attention heads): Multi-Head Attention의 head 수
        self.num_attention_heads = num_attention_heads

        # FFN 내부 차원(intermediate size): Feed Forward Network의 중간 차원
        self.intermediate_size = intermediate_size

        # decoder layer 개수(num hidden layers): DecoderBlock 반복 횟수
        self.num_hidden_layers = num_hidden_layers

        # dropout 비율(hidden dropout probability): 과적합(overfitting) 방지
        self.hidden_dropout_prob = hidden_dropout_prob

        # attention dropout 비율(attention probabilities dropout probability)
        self.attention_probs_dropout_prob = attention_probs_dropout_prob


class PositionWiseFeedForward(nn.Module):
    def __init__(self, hidden_size, intermediate_size, dropout_prob):
        super().__init__()

        # 첫 번째 선형층(first linear layer): hidden size를 intermediate size로 확장
        self.linear_1 = nn.Linear(hidden_size, intermediate_size)

        # 활성화 함수(activation function): 비선형성 추가
        self.activation = nn.GELU()

        # 두 번째 선형층(second linear layer): intermediate size를 hidden size로 축소
        self.linear_2 = nn.Linear(intermediate_size, hidden_size)

        # 드롭아웃(dropout): 과적합(overfitting) 방지
        self.dropout = nn.Dropout(dropout_prob)

    def forward(self, x):
        # Position-wise FFN: 각 위치(position)의 token에 독립적으로 적용
        x = self.linear_1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.linear_2(x)
        return x


class DecoderBlock(nn.Module):
    def __init__(self, config):
        super().__init__()

        # Masked Self-Attention: 미래 token을 보지 못하게 하는 attention
        self.masked_self_attention = nn.MultiheadAttention(
            embed_dim=config.hidden_size,
            num_heads=config.num_attention_heads,
            dropout=config.attention_probs_dropout_prob,
            batch_first=True
        )

        # 첫 번째 정규화(layer normalization): self-attention 출력 안정화
        self.norm_1 = nn.LayerNorm(config.hidden_size)

        # Cross-Attention: encoder output과 decoder input 사이의 관계 학습
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=config.hidden_size,
            num_heads=config.num_attention_heads,
            dropout=config.attention_probs_dropout_prob,
            batch_first=True
        )

        # 두 번째 정규화(layer normalization): cross-attention 출력 안정화
        self.norm_2 = nn.LayerNorm(config.hidden_size)

        # Position-Wise Feed Forward Network: token별 비선형 변환
        self.feed_forward = PositionWiseFeedForward(
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            dropout_prob=config.hidden_dropout_prob
        )

        # 세 번째 정규화(layer normalization): FFN 출력 안정화
        self.norm_3 = nn.LayerNorm(config.hidden_size)

        # 드롭아웃(dropout): residual connection 전에 적용
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, x, encoder_output=None, self_attention_mask=None, cross_attention_mask=None):
        # 잔차 연결(residual connection)을 위해 입력 저장
        residual = x

        # Masked Self-Attention 계산
        self_attention_output, _ = self.masked_self_attention(
            query=x,
            key=x,
            value=x,
            attn_mask=self_attention_mask
        )

        # Add & Norm: residual connection + layer normalization
        x = self.norm_1(residual + self.dropout(self_attention_output))

        if encoder_output is not None:
            # 잔차 연결(residual connection)을 위해 입력 저장
            residual = x

            # Cross-Attention 계산
            cross_attention_output, _ = self.cross_attention(
                query=x,
                key=encoder_output,
                value=encoder_output,
                key_padding_mask=cross_attention_mask
            )

            # Add & Norm: residual connection + layer normalization
            x = self.norm_2(residual + self.dropout(cross_attention_output))

        # 잔차 연결(residual connection)을 위해 FFN 입력 저장
        residual = x

        # Position-Wise Feed Forward 계산
        ff_output = self.feed_forward(x)

        # Add & Norm: residual connection + layer normalization
        x = self.norm_3(residual + self.dropout(ff_output))

        return x


class TransformerDecoder(nn.Module):
    def __init__(self, config):
        super().__init__()

        # 토큰 임베딩(token embedding): token id를 vector로 변환
        self.token_embedding = nn.Embedding(
            config.vocab_size,
            config.hidden_size
        )

        # 위치 임베딩(position embedding): token의 순서 정보 추가
        self.position_embedding = nn.Embedding(
            config.max_position_embeddings,
            config.hidden_size
        )

        # DecoderBlock stack: N개의 DecoderBlock 반복
        self.layers = nn.ModuleList([
            DecoderBlock(config) for _ in range(config.num_hidden_layers)
        ])

        # 드롭아웃(dropout): embedding 이후 적용
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        # 최종 정규화(final layer normalization)
        self.final_norm = nn.LayerNorm(config.hidden_size)

        # 출력층(lm head): hidden vector를 vocab size로 변환
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)

        self.config = config

    def make_causal_mask(self, seq_length, device):
        # Causal mask: 현재 token이 미래 token을 보지 못하게 함
        mask = torch.triu(
            torch.ones(seq_length, seq_length, device=device),
            diagonal=1
        ).bool()
        return mask

    def forward(self, input_ids, encoder_output=None, cross_attention_mask=None):
        batch_size, seq_length = input_ids.size()

        # 위치 번호(position ids) 생성
        position_ids = torch.arange(
            seq_length,
            dtype=torch.long,
            device=input_ids.device
        ).unsqueeze(0).expand(batch_size, seq_length)

        # token embedding과 position embedding을 더함
        x = self.token_embedding(input_ids) + self.position_embedding(position_ids)
        x = self.dropout(x)

        # 미래 token 차단용 causal mask 생성
        self_attention_mask = self.make_causal_mask(seq_length, input_ids.device)

        # DecoderBlock 반복 적용
        for layer in self.layers:
            x = layer(
                x,
                encoder_output=encoder_output,
                self_attention_mask=self_attention_mask,
                cross_attention_mask=cross_attention_mask
            )

        x = self.final_norm(x)

        # logits: 각 token 위치에서 vocab 전체에 대한 점수
        logits = self.lm_head(x)

        return logits


if __name__ == "__main__":
    config = CustomDecoderConfig()
    model = TransformerDecoder(config)

    input_ids = torch.randint(0, config.vocab_size, (2, 10))
    encoder_output = torch.randn(2, 10, config.hidden_size)

    logits = model(input_ids, encoder_output=encoder_output)

    print("Input shape:", input_ids.shape)
    print("Logits shape:", logits.shape)