Transformer Architecture Implementations

Standard Transformer

Components:

1. Multi-Head Self-Attention (MultiHeadSelfAttention class)
   • Scaled dot-product attention
   • Multiple attention heads
   python

   1# Key implementation snippet 2class MultiHeadSelfAttention(nn.Module): 3 def __init__(self, d_model, num_heads, dropout=0.1): 4 # Initialize projections for queries, keys, and values 5 self.qkv_proj = nn.Linear(d_model, 3 * d_model) 6 self.out_proj = nn.Linear(d_model, d_model) 7 self.scale = self.head_dim ** -0.5
2. Position-wise Feed-Forward Network (FeedForward class)
   • Two linear transformations with GELU activation
   python

   1# Key implementation snippet 2class FeedForward(nn.Module): 3 def __init__(self, d_model, d_ff, dropout=0.1): 4 self.linear1 = nn.Linear(d_model, d_ff) 5 self.linear2 = nn.Linear(d_ff, d_model) 6 self.activation = nn.GELU()
3. Encoder Block (EncoderBlock class)
   • Self-attention + Feed-forward with residual connections and layer norm
   python

   1# Key implementation snippet 2class EncoderBlock(nn.Module): 3 def __init__(self, d_model, num_heads, d_ff, dropout=0.1): 4 self.self_attn = MultiHeadSelfAttention(d_model, num_heads, dropout) 5 self.norm1 = nn.LayerNorm(d_model) 6 self.ff = FeedForward(d_model, d_ff, dropout) 7 self.norm2 = nn.LayerNorm(d_model)
4. Full Transformer (StandardTransformer class)
   • Stack of encoder blocks
   python

   1# Key implementation snippet 2class StandardTransformer(nn.Module): 3 def __init__(self, d_model, num_heads, d_ff, num_layers, dropout=0.1): 4 self.layers = nn.ModuleList([ 5 EncoderBlock(d_model, num_heads, d_ff, dropout) 6 for _ in range(num_layers) 7 ])

SLAM v1 Architecture

Architecture Breakdown:

1. Initial Full Self-Attention Layer (SLAMv1.initial_block)
   • Purpose: Capture global context before entering EF cycles
   • Components:
     • Multi-Head Self-Attention (MHSA)
     • Layer Normalization
     • Feed-Forward Network (FFN)
     • Layer Normalization
   python

   1# Key implementation snippet 2class SLAMv1(nn.Module): 3 def __init__(self, d_model, num_heads, d_ff, ef_cycles, dropout=0.1): 4 # Initial full MHSA layer 5 self.initial_block = EncoderBlock(d_model, num_heads, d_ff, dropout)
2. Encoder Fusion (EF) Cycles – Main Body (SLAMv1.ef_blocks_A and SLAMv1.ef_blocks_B)
   • Each EF cycle has two encoder blocks processing overlapping chunks
   python

   1# Key implementation snippet 2# EF Cycle blocks 3self.ef_blocks_A = nn.ModuleList([ 4 EncoderBlock(d_model, num_heads, d_ff, dropout) 5 for _ in range(ef_cycles) 6]) 7 8self.ef_blocks_B = nn.ModuleList([ 9 EncoderBlock(d_model, num_heads, d_ff, dropout) 10 for _ in range(ef_cycles) 11])
   • Segment Processing (SLAMv1.forward method)
   python

   1# Key implementation snippet 2# Block A processes first 60% and last 60% segments 3segment_A1 = x[:, :first_60_percent, :] 4segment_A2 = x[:, start_40_percent:, :] 5segment_A = torch.cat([segment_A1, segment_A2], dim=1) 6 7# Block B processes last 60% and first 60% segments (order reversed from A) 8segment_B1 = x[:, start_40_percent:, :] 9segment_B2 = x[:, :first_60_percent, :] 10segment_B = torch.cat([segment_B1, segment_B2], dim=1)
   • Fusion Logic (overlapping region handling)
   python

   1# Key implementation snippet 2# Overlapping region: [start_40_percent, first_60_percent) 3x_new[:, overlap_start:overlap_end, :] += processed_A1[:, overlap_start:overlap_end, :] 4counts[:, overlap_start:overlap_end, :] += 1 5 6x_new[:, overlap_start:overlap_end, :] += processed_B2[:, overlap_start:overlap_end, :] 7counts[:, overlap_start:overlap_end, :] += 1 8 9# Average the overlapping regions 10x = x_new / counts
3. Final Refinement Layer (SLAMv1.final_block)
   • Purpose: Reintroduce full-sequence attention to integrate context
   python

   1# Key implementation snippet 2# Final refinement block 3self.final_block = EncoderBlock(d_model, num_heads, d_ff, dropout) 4 5# In forward method: 6x = self.final_block(x, mask)

SLAM v2 Architecture

Architecture Breakdown:

1. Level 1: Initial Global Attention (SLAMv2.initial_block)
   • MoE Attention Layer (20 experts = 20 distinct QKV heads)
   • Normalization
   • Feed Forward Network (FFN)
   • Normalization
   python

   1# Key implementation snippet 2class SLAMv2(nn.Module): 3 def __init__(self, d_model, num_experts=20, d_ff=2048, ef_cycles=2, dropout=0.1): 4 # Level 1: Initial Global Attention 5 self.initial_block = MoEEncoderBlock(d_model, num_experts, d_ff, dropout)
2. Level 2: Encoder Fusion (EF) Cycles (SLAMv2.ef_blocks)
   • 4 parallel encoder blocks, each with MoE attention
   python

   1# Key implementation snippet 2# Level 2: Encoder Fusion (EF) Cycles 3# 4 parallel encoder blocks for each EF cycle 4self.ef_blocks = nn.ModuleList([ 5 nn.ModuleList([ 6 MoEEncoderBlock(d_model, num_experts, d_ff, dropout) 7 for _ in range(4) # 4 parallel blocks 8 ]) 9 for _ in range(ef_cycles) 10])
   • Segment Boundaries (token ranges for each block)
   python

   1# Key implementation snippet 2segment_boundaries = [ 3 (0, min(int(seq_len * 0.45), seq_len)), # Block 1: tokens 0-45% 4 (max(int(seq_len * 0.25), 0), min(int(seq_len * 0.65), seq_len)), # Block 2: tokens 25-65% 5 (max(int(seq_len * 0.45), 0), min(int(seq_len * 0.85), seq_len)), # Block 3: tokens 45-85% 6 (max(int(seq_len * 0.65), 0), seq_len) # Block 4: tokens 65-100% 7] 8 9# Block 4 with wraparound 10segment_boundaries[3] = ( 11 max(int(seq_len * 0.65), 0), # Start at 65% 12 seq_len + wrap_size # End at 100% + 5% wraparound 13)
3. Level 3: Final Aggregation (SLAMv2.final_block)
   • Final MoE Attention Layer with normalization and FFN
   python

   1# Key implementation snippet 2# Level 3: Final Aggregation 3self.final_block = MoEEncoderBlock(d_model, num_experts, d_ff, dropout) 4 5# In forward method: 6x = self.final_block(x, mask)

Mixture of Experts (MoE) Implementation

Architecture Breakdown:

1. Router Network (MoEAttention.router)
   • Determines which expert(s) each token should be processed by
   python

   1# Key implementation snippet 2class MoEAttention(nn.Module): 3 def __init__(self, d_model, num_experts=20, dropout=0.1): 4 # Router network to determine which expert to use 5 self.router = nn.Linear(d_model, num_experts) 6 self.router_dropout = nn.Dropout(0.1) # Dropout for router logits
2. Expert Networks (MoEAttention.experts)
   • Each expert is a specialized QKV attention head
   python

   1# Key implementation snippet 2# Each expert is a QKV head 3self.experts = nn.ModuleList([ 4 nn.Linear(d_model, 3 * self.head_dim) for _ in range(num_experts) 5])
3. Token Routing (in forward method)
   • Each token is processed by the most relevant expert(s)
   python

   1# Key implementation snippet 2# Calculate routing probabilities with dropout for regularization 3router_logits = self.router(x) # [batch_size, seq_len, num_experts] 4router_logits = self.router_dropout(router_logits) 5routing_weights = torch.softmax(router_logits, dim=-1)
4. MoE Encoder Block (MoEEncoderBlock class)
   • Combines MoE attention with feed-forward network
   python

   1# Key implementation snippet 2class MoEEncoderBlock(nn.Module): 3 def __init__(self, d_model, num_experts, d_ff, dropout=0.1): 4 self.self_attn = MoEAttention(d_model, num_experts, dropout) 5 self.norm1 = nn.LayerNorm(d_model) 6 self.ff = FeedForward(d_model, d_ff, dropout) 7 self.norm2 = nn.LayerNorm(d_model)

DeepSeek Transformer

Architecture Breakdown:

1. Rotary Position Embeddings (RoPE) (RotaryEmbedding class)
   • Applied to queries and keys to encode relative positions
   python

   1# Key implementation snippet 2class RotaryEmbedding(nn.Module): 3 def __init__(self, dim, max_seq_len=2048, base=10000): 4 # Create the frequency base 5 inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) 6 self.register_buffer("inv_freq", inv_freq)
2. Grouped Query Attention (GQA) (TrueGroupedQueryAttention class)
   • Reduces computational cost by sharing queries across multiple heads
   python

   1# Key implementation snippet 2class TrueGroupedQueryAttention(nn.Module): 3 def __init__(self, d_model, num_heads, num_query_groups=1, dropout=0.1, max_seq_len=2048): 4 # True GQA: Only project queries for each group (fewer projections) 5 self.q_proj = nn.Linear(d_model, self.head_dim * num_query_groups) 6 7 # Keys and values are still full-sized (one per head) 8 self.k_proj = nn.Linear(d_model, d_model) 9 self.v_proj = nn.Linear(d_model, d_model)
3. DeepSeek Block (DeepSeekBlock class)
   • Pre-normalization + GQA + SwiGLU feed-forward
   python

   1# Key implementation snippet 2class DeepSeekBlock(nn.Module): 3 def __init__(self, d_model, num_heads, d_ff, num_query_groups=1, dropout=0.1, max_seq_len=2048): 4 self.pre_norm1 = RMSNorm(d_model) 5 self.self_attn = TrueGroupedQueryAttention(d_model, num_heads, num_query_groups, dropout, max_seq_len) 6 self.pre_norm2 = RMSNorm(d_model) 7 self.ff = SwiGLU(d_model, d_ff, dropout)
4. Full DeepSeek Transformer (DeepSeekTransformer class)
   • Stack of DeepSeek blocks with final normalization
   python

   1# Key implementation snippet 2class DeepSeekTransformer(nn.Module): 3 def __init__(self, d_model, num_heads, d_ff, num_layers, num_query_groups=1, dropout=0.1, max_seq_len=2048): 4 self.layers = nn.ModuleList([ 5 DeepSeekBlock(d_model, num_heads, d_ff, num_query_groups, dropout, max_seq_len) 6 for _ in range(num_layers) 7 ]) 8 self.final_norm = RMSNorm(d_model)

DeBERTa Architecture

Architecture Breakdown:

1. Disentangled Self-Attention (DisentangledSelfAttention class)
   • Separates content-to-content, content-to-position, and position-to-content attention
   python

   1# Key implementation snippet 2class DisentangledSelfAttention(nn.Module): 3 def __init__(self, d_model, num_heads, dropout=0.1, max_position_embeddings=512): 4 # Content projections with explicit bias terms 5 self.q_proj = nn.Linear(d_model, d_model, bias=True) 6 self.k_proj = nn.Linear(d_model, d_model, bias=True) 7 self.v_proj = nn.Linear(d_model, d_model, bias=True) 8 9 # Position projections 10 self.pos_k_proj = nn.Linear(d_model, d_model, bias=False) 11 self.pos_q_proj = nn.Linear(d_model, d_model, bias=False)
2. Relative Position Embeddings
   • Explicitly models relative token positions
   python

   1# Key implementation snippet 2# Relative position embeddings table (2*max_len - 1) positions 3self.pos_embeddings = nn.Parameter( 4 torch.zeros(2 * max_position_embeddings - 1, d_model) 5) 6nn.init.normal_(self.pos_embeddings, mean=0, std=0.02)
3. DeBERTa Encoder Block (DeBERTaEncoderBlock class)
   • Combines disentangled attention with feed-forward network
   python

   1# Key implementation snippet 2class DeBERTaEncoderBlock(nn.Module): 3 def __init__(self, d_model, num_heads, d_ff, dropout=0.1, max_position_embeddings=512): 4 self.self_attn = DisentangledSelfAttention(d_model, num_heads, dropout, max_position_embeddings) 5 # Use RMSNorm instead of LayerNorm for better performance 6 self.norm1 = RMSNorm(d_model) 7 # Use SwiGLU instead of standard FeedForward for better performance 8 self.ff = SwiGLU(d_model, d_ff, dropout) 9 self.norm2 = RMSNorm(d_model)
4. Full DeBERTa Encoder (DeBERTaEncoder class)
   • Stack of DeBERTa blocks with final normalization
   python

   1# Key implementation snippet 2class DeBERTaEncoder(nn.Module): 3 def __init__(self, d_model, num_heads, d_ff, num_layers, dropout=0.1, max_position_embeddings=512): 4 # Encoder layers 5 self.layers = nn.ModuleList([ 6 DeBERTaEncoderBlock(d_model, num_heads, d_ff, dropout, max_position_embeddings) 7 for _ in range(num_layers) 8 ]) 9 10 # Final normalization - use RMSNorm for consistency 11 self.norm = RMSNorm(d_model)

Classifier Wrapper