gated_attention_model.py

"""
Gated Attention for Large Language Models
Implementation based on "Gated Attention for Large Language Models: 
Non-linearity, Sparsity, and Attention-Sink-Free" (Qiu et al., 2025)

This module implements the core gated attention mechanism.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple
import math


class GatedAttention(nn.Module):
    """
    Multi-head attention with optional gating mechanisms.
    
    Supports various gating positions:
    - G1: After SDPA output (most effective)
    - G2: After value projection
    - G3: After key projection
    - G4: After query projection
    - G5: After dense output layer
    """
    
    def __init__(
        self,
        d_model: int,
        n_heads: int,
        n_kv_heads: Optional[int] = None,
        head_dim: Optional[int] = None,
        dropout: float = 0.0,
        bias: bool = False,
        gate_position: str = "G1",  # G1, G2, G3, G4, G5, or None
        gate_type: str = "elementwise",  # "elementwise" or "headwise"
        gate_activation: str = "sigmoid",  # "sigmoid" or "silu"
        gate_shared: bool = False,  # Head-specific or head-shared
        gate_multiplicative: bool = True,  # Multiplicative vs additive
    ):
        super().__init__()
        
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads
        self.head_dim = head_dim if head_dim is not None else d_model // n_heads
        self.dropout = dropout
        self.gate_position = gate_position
        self.gate_type = gate_type
        self.gate_activation = gate_activation
        self.gate_shared = gate_shared
        self.gate_multiplicative = gate_multiplicative
        
        # Ensure GQA compatibility
        assert self.n_heads % self.n_kv_heads == 0, \
            "n_heads must be divisible by n_kv_heads"
        self.n_rep = self.n_heads // self.n_kv_heads
        
        # QKV projections
        self.q_proj = nn.Linear(d_model, n_heads * self.head_dim, bias=bias)
        self.k_proj = nn.Linear(d_model, self.n_kv_heads * self.head_dim, bias=bias)
        self.v_proj = nn.Linear(d_model, self.n_kv_heads * self.head_dim, bias=bias)
        
        # Output projection
        self.o_proj = nn.Linear(n_heads * self.head_dim, d_model, bias=bias)
        
        # Dropout
        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)
        
        # Initialize gating parameters
        self._init_gating()
        
    def _init_gating(self):
        """Initialize gating parameters based on configuration."""
        if self.gate_position is None or self.gate_position == "none":
            self.gate = None
            return
            
        # Determine gate input/output dimensions
        if self.gate_position == "G1":  # After SDPA
            if self.gate_type == "elementwise":
                if self.gate_shared:
                    gate_dim = self.head_dim
                else:
                    gate_dim = self.n_heads * self.head_dim
            else:  # headwise
                gate_dim = self.n_heads
                
        elif self.gate_position == "G2":  # After value projection
            if self.gate_type == "elementwise":
                if self.gate_shared:
                    gate_dim = self.head_dim
                else:
                    gate_dim = self.n_kv_heads * self.head_dim
            else:  # headwise
                gate_dim = self.n_kv_heads
                
        elif self.gate_position in ["G3", "G4"]:  # After K or Q projection
            if self.gate_type == "elementwise":
                target_heads = self.n_kv_heads if self.gate_position == "G3" else self.n_heads
                if self.gate_shared:
                    gate_dim = self.head_dim
                else:
                    gate_dim = target_heads * self.head_dim
            else:  # headwise
                gate_dim = self.n_kv_heads if self.gate_position == "G3" else self.n_heads
                
        elif self.gate_position == "G5":  # After output projection
            if self.gate_type == "elementwise":
                gate_dim = self.d_model
            else:  # headwise
                gate_dim = 1
        else:
            raise ValueError(f"Unknown gate position: {self.gate_position}")
        
        # Create gate projection
        self.gate = nn.Linear(self.d_model, gate_dim, bias=False)
        
        # Initialize gate weights to produce small values initially
        with torch.no_grad():
            self.gate.weight.normal_(mean=0.0, std=0.02)
    
    def _apply_gate(
        self, 
        x: torch.Tensor, 
        hidden_states: torch.Tensor,
        position: str
    ) -> torch.Tensor:
        """Apply gating to input tensor."""
        if self.gate is None or self.gate_position != position:
            return x
        
        # Compute gate scores
        gate_scores = self.gate(hidden_states)
        
        # Apply activation
        if self.gate_activation == "sigmoid":
            gate_scores = torch.sigmoid(gate_scores)
        elif self.gate_activation == "silu":
            gate_scores = F.silu(gate_scores)
        else:
            raise ValueError(f"Unknown activation: {self.gate_activation}")
        
        # Reshape for broadcasting if needed
        if self.gate_position == "G1" and self.gate_type == "elementwise":
            if self.gate_shared:
                # gate_scores: [batch, seq, head_dim]
                # x: [batch, seq, n_heads, head_dim]
                gate_scores = gate_scores.unsqueeze(2)
            else:
                # gate_scores: [batch, seq, n_heads * head_dim]
                # Reshape to [batch, seq, n_heads, head_dim]
                batch, seq, _ = gate_scores.shape
                gate_scores = gate_scores.view(batch, seq, self.n_heads, self.head_dim)
                
        elif self.gate_position == "G1" and self.gate_type == "headwise":
            # gate_scores: [batch, seq, n_heads]
            # Expand to [batch, seq, n_heads, 1]
            gate_scores = gate_scores.unsqueeze(-1)
        
        # Apply gate (multiplicative or additive)
        if self.gate_multiplicative:
            return x * gate_scores
        else:
            return x + gate_scores
    
    def repeat_kv(self, x: torch.Tensor) -> torch.Tensor:
        """Repeat key/value tensors for GQA."""
        if self.n_rep == 1:
            return x
        batch, seq, n_kv_heads, head_dim = x.shape
        return x.repeat_interleave(self.n_rep, dim=2)
    
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        batch_size, seq_length, _ = hidden_states.shape
        
        # QKV projections
        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)
        
        # Reshape to [batch, seq, n_heads, head_dim]
        query_states = query_states.view(
            batch_size, seq_length, self.n_heads, self.head_dim
        )
        key_states = key_states.view(
            batch_size, seq_length, self.n_kv_heads, self.head_dim
        )
        value_states = value_states.view(
            batch_size, seq_length, self.n_kv_heads, self.head_dim
        )
        
        # Apply gating at G4, G3, G2 positions
        query_states = self._apply_gate(query_states, hidden_states, "G4")
        key_states = self._apply_gate(key_states, hidden_states, "G3")
        value_states = self._apply_gate(value_states, hidden_states, "G2")
        
        # Handle KV cache
        if past_key_value is not None:
            key_states = torch.cat([past_key_value[0], key_states], dim=1)
            value_states = torch.cat([past_key_value[1], value_states], dim=1)
        
        if use_cache:
            past_key_value = (key_states, value_states)
        else:
            past_key_value = None
        
        # Repeat KV for GQA
        key_states = self.repeat_kv(key_states)
        value_states = self.repeat_kv(value_states)
        
        # Transpose for attention: [batch, n_heads, seq, head_dim]
        query_states = query_states.transpose(1, 2)
        key_states = key_states.transpose(1, 2)
        value_states = value_states.transpose(1, 2)
        
        # Scaled dot-product attention
        attn_weights = torch.matmul(
            query_states, key_states.transpose(-2, -1)
        ) / math.sqrt(self.head_dim)
        
        if attention_mask is not None:
            attn_weights = attn_weights + attention_mask
        
        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
            query_states.dtype
        )
        attn_weights = self.attn_dropout(attn_weights)
        
        # Compute attention output
        attn_output = torch.matmul(attn_weights, value_states)
        
        # Transpose back: [batch, seq, n_heads, head_dim]
        attn_output = attn_output.transpose(1, 2).contiguous()
        
        # Apply gating at G1 position (SDPA output)
        attn_output = self._apply_gate(attn_output, hidden_states, "G1")
        
        # Reshape to [batch, seq, n_heads * head_dim]
        attn_output = attn_output.view(batch_size, seq_length, -1)
        
        # Output projection
        attn_output = self.o_proj(attn_output)
        
        # Apply gating at G5 position
        attn_output = self._apply_gate(attn_output, hidden_states, "G5")
        
        attn_output = self.resid_dropout(attn_output)
        
        outputs = (attn_output,)
        if output_attentions:
            outputs += (attn_weights,)
        if use_cache:
            outputs += (past_key_value,)
        
        return outputs


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization."""
    
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x * norm * self.weight


class FeedForward(nn.Module):
    """
    SwiGLU Feed-Forward Network.
    FFN(x) = (Swish(xW1) ⊙ xW3)W2
    """
    
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        dropout: float = 0.0,
        bias: bool = False,
    ):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=bias)
        self.w2 = nn.Linear(d_ff, d_model, bias=bias)
        self.w3 = nn.Linear(d_model, d_ff, bias=bias)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))


class TransformerBlock(nn.Module):
    """Transformer block with gated attention and SwiGLU FFN."""
    
    def __init__(
        self,
        d_model: int,
        n_heads: int,
        n_kv_heads: Optional[int] = None,
        d_ff: Optional[int] = None,
        dropout: float = 0.0,
        norm_eps: float = 1e-6,
        gate_position: str = "G1",
        gate_type: str = "elementwise",
        gate_activation: str = "sigmoid",
        gate_shared: bool = False,
        sandwich_norm: bool = False,
    ):
        super().__init__()
        
        self.attention = GatedAttention(
            d_model=d_model,
            n_heads=n_heads,
            n_kv_heads=n_kv_heads,
            dropout=dropout,
            gate_position=gate_position,
            gate_type=gate_type,
            gate_activation=gate_activation,
            gate_shared=gate_shared,
        )
        
        d_ff = d_ff or 4 * d_model
        self.feed_forward = FeedForward(
            d_model=d_model,
            d_ff=d_ff,
            dropout=dropout,
        )
        
        self.attention_norm = RMSNorm(d_model, eps=norm_eps)
        self.ffn_norm = RMSNorm(d_model, eps=norm_eps)
        
        # Sandwich norm (post-attention/FFN normalization)
        self.sandwich_norm = sandwich_norm
        if sandwich_norm:
            self.post_attention_norm = RMSNorm(d_model, eps=norm_eps)
            self.post_ffn_norm = RMSNorm(d_model, eps=norm_eps)
    
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
    ):
        # Attention with pre-norm
        residual = hidden_states
        hidden_states = self.attention_norm(hidden_states)
        
        attn_outputs = self.attention(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
        )
        
        hidden_states = attn_outputs[0]
        
        # Optional sandwich norm
        if self.sandwich_norm:
            hidden_states = self.post_attention_norm(hidden_states)
        
        hidden_states = residual + hidden_states
        
        # FFN with pre-norm
        residual = hidden_states
        hidden_states = self.ffn_norm(hidden_states)
        hidden_states = self.feed_forward(hidden_states)
        
        # Optional sandwich norm
        if self.sandwich_norm:
            hidden_states = self.post_ffn_norm(hidden_states)
        
        hidden_states = residual + hidden_states
        
        outputs = (hidden_states,)
        if output_attentions:
            outputs += (attn_outputs[1],)
        if use_cache:
            outputs += (attn_outputs[-1],)
        
        return outputs