class TransformerEncoderBlock(nn.Module):
    def __init__(self, d_model=512, num_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.attention = nn.MultiheadAttention(d_model, num_heads,
                                               dropout=dropout, batch_first=True)
        self.ffn = FeedForward(d_model, d_ff, dropout)

        # Pre-LN (modern default) — norm before sublayer
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, src_key_padding_mask=None):
        # Pre-LN: norm -> sublayer -> add residual
        # Multi-Head Self-Attention
        normed = self.norm1(x)
        attn_out, _ = self.attention(normed, normed, normed,
                                     key_padding_mask=src_key_padding_mask)
        x = x + self.dropout(attn_out)  # residual

        # Feed-Forward
        normed = self.norm2(x)
        x = x + self.dropout(self.ffn(normed))  # residual

        return x  # [batch, seq_len, d_model]

Transformers — ארכיטקטורה שמניעה את כל ה-AI המודרני

Transformers — ארכיטקטורה שמניעה את כל ה-AI המודרני