Initial commit

crates/training/rtx-transformers/src/layers/linear_attention.rs

@@ -0,0 +1,393 @@
+//! Linear Attention Implementation
+//!
+//! Reduces O(n²) complexity to O(n) by using feature map approximations φ(Q)φ(K)ᵀ
+//! 
+//! Key features:
+//! - Feature map approximation for attention computation
+//! - Causal masking support for autoregressive models
+//! - Memory efficient computation avoiding n×n attention matrices
+//! - Compatible with existing transformer infrastructure
+
+use crate::{Result, TransformerError};
+use rtx_tensor::{Tensor, Device, DType};
+use crate::layers::{Layer, LayerNorm};
+
+/// Feature map types for linear attention
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum FeatureMapType {
+    /// ELU activation: max(0, x) + 1
+    Elu,
+    /// ReLU activation: max(0, x)
+    Relu,
+    /// Softmax normalization
+    Softmax,
+    /// Random Fourier features
+    Fourier,
+}
+
+/// Configuration for Linear Attention
+#[derive(Debug, Clone)]
+pub struct LinearAttentionConfig {
+    pub d_model: usize,
+    pub num_heads: usize,
+    pub head_dim: usize,
+    pub feature_map_dim: usize,
+    pub causal: bool,
+    pub feature_map_type: FeatureMapType,
+    pub epsilon: f32,
+}
+
+impl LinearAttentionConfig {
+    /// Create a new linear attention configuration
+    pub fn new(d_model: usize, num_heads: usize) -> Self {
+        let head_dim = d_model / num_heads;
+        
+        Self {
+            d_model,
+            num_heads,
+            head_dim,
+            feature_map_dim: head_dim, // default same as head_dim
+            causal: false,
+            feature_map_type: FeatureMapType::Elu,
+            epsilon: 1e-5,
+        }
+    }
+}
+
+/// Feature map for linear attention approximation
+#[derive(Debug)]
+pub struct FeatureMap {
+    feature_type: FeatureMapType,
+    input_dim: usize,
+    feature_dim: usize,
+    // For Fourier features
+    frequencies: Option<Tensor>,
+    device: Device,
+}
+
+impl FeatureMap {
+    /// Create a new feature map
+    pub fn new(input_dim: usize, feature_dim: usize, feature_type: FeatureMapType, device: &Device) -> Result<Self> {
+        let frequencies = if feature_type == FeatureMapType::Fourier {
+            // Random Fourier features
+            let freq = Tensor::randn(&[input_dim, feature_dim], DType::F32, device)?;
+            Some(freq)
+        } else {
+            None
+        };
+        
+        Ok(Self {
+            feature_type,
+            input_dim,
+            feature_dim,
+            frequencies,
+            device: device.clone(),
+        })
+    }
+    
+    /// Apply feature map transformation
+    pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
+        match self.feature_type {
+            FeatureMapType::Elu => {
+                // φ(x) = ELU(x) + 1 for non-negative features
+                let elu_output = input.elu(1.0)?;
+                let output = (&elu_output + 1.0)?;
+                Ok(output)
+            },
+            FeatureMapType::Relu => {
+                // φ(x) = ReLU(x)
+                input.relu()
+            },
+            FeatureMapType::Softmax => {
+                // φ(x) = softmax(x) (along last dimension)
+                input.softmax(-1)
+            },
+            FeatureMapType::Fourier => {
+                // φ(x) = [cos(Wx), sin(Wx)] where W is random frequencies
+                if let Some(ref frequencies) = self.frequencies {
+                    let projected = input.matmul(frequencies)?;
+                    let cos_part = projected.cos()?;
+                    let sin_part = projected.sin()?;
+                    
+                    // Concatenate cos and sin parts
+                    let output = Tensor::cat(&[cos_part, sin_part], -1)?;
+                    
+                    // Scale by sqrt(feature_dim) for normalization
+                    let scale = (self.feature_dim as f32).sqrt().recip();
+                    let output = (&output * scale)?;
+                    Ok(output)
+                } else {
+                    return Err(TransformerError::InvalidInput("Fourier frequencies not initialized".to_string()));
+                }
+            }
+        }
+    }
+}
+
+/// Fourier feature map implementation
+#[derive(Debug)]
+pub struct FourierFeatureMap {
+    frequencies: Tensor,
+    scale: f32,
+    device: Device,
+}
+
+impl FourierFeatureMap {
+    /// Create a new Fourier feature map
+    pub fn new(input_dim: usize, feature_dim: usize, scale: f32, device: &Device) -> Result<Self> {
+        // Random frequencies sampled from Gaussian distribution
+        let frequencies = Tensor::randn(&[input_dim, feature_dim / 2], DType::F32, device)? * scale;
+        
+        Ok(Self {
+            frequencies,
+            scale,
+            device: device.clone(),
+        })
+    }
+    
+    /// Forward pass through Fourier feature map
+    pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
+        let projected = input.matmul(&self.frequencies)?;
+        let cos_part = projected.cos()?;
+        let sin_part = projected.sin()?;
+        
+        // Concatenate and normalize
+        let output = Tensor::cat(&[cos_part, sin_part], -1)?;
+        let scale_factor = (self.frequencies.dims()[1] as f32 * 2.0).sqrt().recip();
+        let output = (&output * scale_factor)?;
+        
+        Ok(output)
+    }
+}
+
+/// Core linear attention mechanism
+#[derive(Debug)]
+pub struct LinearAttention {
+    q_proj: Tensor,
+    k_proj: Tensor,
+    v_proj: Tensor,
+    o_proj: Tensor,
+    feature_map: FeatureMap,
+    config: LinearAttentionConfig,
+    device: Device,
+}
+
+impl LinearAttention {
+    /// Create a new linear attention layer
+    pub fn new(config: &LinearAttentionConfig, device: &Device) -> Result<Self> {
+        // Initialize projection matrices
+        let scale = (config.d_model as f32).sqrt().recip();
+        
+        let q_proj = Tensor::randn(&[config.d_model, config.d_model], DType::F32, device)? * scale;
+        let k_proj = Tensor::randn(&[config.d_model, config.d_model], DType::F32, device)? * scale;
+        let v_proj = Tensor::randn(&[config.d_model, config.d_model], DType::F32, device)? * scale;
+        let o_proj = Tensor::randn(&[config.d_model, config.d_model], DType::F32, device)? * scale;
+        
+        let feature_map = FeatureMap::new(
+            config.head_dim,
+            config.feature_map_dim,
+            config.feature_map_type,
+            device,
+        )?;
+        
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            feature_map,
+            config: config.clone(),
+            device: device.clone(),
+        })
+    }
+    
+    /// Forward pass with self-attention
+    pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
+        self.forward_with_kv(input, input, input)
+    }
+    
+    /// Forward pass with separate K, V inputs (cross-attention)
+    pub fn forward_with_kv(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
+        let input_shape = q.dims();
+        let batch_size = input_shape[0];
+        let seq_len = input_shape[1];
+        
+        // Project inputs to Q, K, V
+        let q_proj = q.matmul(&self.q_proj)?;
+        let k_proj = k.matmul(&self.k_proj)?;
+        let v_proj = v.matmul(&self.v_proj)?;
+        
+        // Reshape for multi-head attention
+        let q_proj = q_proj.reshape(&[batch_size, seq_len, self.config.num_heads, self.config.head_dim])?
+                          .transpose(1, 2)?; // [B, H, T, D]
+        let k_proj = k_proj.reshape(&[batch_size, seq_len, self.config.num_heads, self.config.head_dim])?
+                          .transpose(1, 2)?; // [B, H, T, D]
+        let v_proj = v_proj.reshape(&[batch_size, seq_len, self.config.num_heads, self.config.head_dim])?
+                          .transpose(1, 2)?; // [B, H, T, D]
+        
+        // Apply feature maps: φ(Q) and φ(K)
+        let phi_q = self.apply_feature_map_to_heads(&q_proj)?; // [B, H, T, F]
+        let phi_k = self.apply_feature_map_to_heads(&k_proj)?; // [B, H, T, F]
+        
+        // Compute linear attention
+        let output = if self.config.causal {
+            self.causal_linear_attention(&phi_q, &phi_k, &v_proj)?
+        } else {
+            self.non_causal_linear_attention(&phi_q, &phi_k, &v_proj)?
+        };
+        
+        // Reshape back and apply output projection
+        let output = output.transpose(1, 2)?.reshape(&[batch_size, seq_len, self.config.d_model])?;
+        let output = output.matmul(&self.o_proj)?;
+        
+        Ok(output)
+    }
+    
+    /// Apply feature map to multi-head tensor
+    fn apply_feature_map_to_heads(&self, input: &Tensor) -> Result<Tensor> {
+        let input_shape = input.dims();
+        let batch_size = input_shape[0];
+        let num_heads = input_shape[1];
+        let seq_len = input_shape[2];
+        let head_dim = input_shape[3];
+        
+        // Reshape to apply feature map: [B, H, T, D] -> [B*H*T, D]
+        let reshaped = input.reshape(&[batch_size * num_heads * seq_len, head_dim])?;
+        let features = self.feature_map.forward(&reshaped)?;
+        
+        // Reshape back: [B*H*T, F] -> [B, H, T, F]
+        let feature_dim = features.dims()[1];
+        let output = features.reshape(&[batch_size, num_heads, seq_len, feature_dim])?;
+        
+        Ok(output)
+    }
+    
+    /// Non-causal linear attention: O = φ(Q)(φ(K)ᵀV)
+    fn non_causal_linear_attention(&self, phi_q: &Tensor, phi_k: &Tensor, v: &Tensor) -> Result<Tensor> {
+        // φ(K)ᵀV: [B, H, F, T] × [B, H, T, D] -> [B, H, F, D]
+        let kv = phi_k.transpose(-1, -2)?.matmul(v)?;
+        
+        // φ(Q)(φ(K)ᵀV): [B, H, T, F] × [B, H, F, D] -> [B, H, T, D]
+        let output = phi_q.matmul(&kv)?;
+        
+        Ok(output)
+    }
+    
+    /// Causal linear attention with cumulative computation
+    fn causal_linear_attention(&self, phi_q: &Tensor, phi_k: &Tensor, v: &Tensor) -> Result<Tensor> {
+        let input_shape = phi_q.dims();
+        let batch_size = input_shape[0];
+        let num_heads = input_shape[1];
+        let seq_len = input_shape[2];
+        let head_dim = v.dims()[3];
+        
+        let mut outputs = Vec::new();
+        let mut cumulative_kv = Tensor::zeros(&[batch_size, num_heads, self.config.feature_map_dim, head_dim], &self.device)?;
+        
+        // Process each position causally
+        for t in 0..seq_len {
+            let phi_k_t = phi_k.narrow(2, t, 1)?; // [B, H, 1, F]
+            let v_t = v.narrow(2, t, 1)?; // [B, H, 1, D]
+            
+            // Update cumulative K*V
+            let kv_t = phi_k_t.transpose(-1, -2)?.matmul(&v_t)?; // [B, H, F, D]
+            cumulative_kv = (&cumulative_kv + &kv_t)?;
+            
+            // Compute output for position t
+            let phi_q_t = phi_q.narrow(2, t, 1)?; // [B, H, 1, F]
+            let output_t = phi_q_t.matmul(&cumulative_kv)?; // [B, H, 1, D]
+            
+            outputs.push(output_t);
+        }
+        
+        // Concatenate outputs
+        let output = Tensor::cat(&outputs, 2)?;
+        Ok(output)
+    }
+}
+
+/// Linear attention layer with normalization and residual connections
+#[derive(Debug)]
+pub struct LinearAttentionLayer {
+    attention: LinearAttention,
+    norm: LayerNorm,
+    config: LinearAttentionConfig,
+    device: Device,
+}
+
+impl LinearAttentionLayer {
+    /// Create a new linear attention layer
+    pub fn new(config: &LinearAttentionConfig, device: &Device) -> Result<Self> {
+        let attention = LinearAttention::new(config, device)?;
+        let norm = LayerNorm::new(config.d_model, config.epsilon, device)?;
+        
+        Ok(Self {
+            attention,
+            norm,
+            config: config.clone(),
+            device: device.clone(),
+        })
+    }
+    
+    /// Forward pass with pre-normalization and residual connection
+    pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
+        // Pre-normalization
+        let normalized = self.norm.forward(input)?;
+        
+        // Linear attention
+        let attn_output = self.attention.forward(&normalized)?;
+        
+        // Residual connection
+        let output = (input + &attn_output)?;
+        
+        Ok(output)
+    }
+}
+
+impl Layer for LinearAttention {
+    fn forward(&self, input: &Tensor) -> Result<Tensor> {
+        self.forward(input)
+    }
+    
+    fn layer_type(&self) -> &'static str {
+        "LinearAttention"
+    }
+    
+    fn device(&self) -> &Device {
+        &self.device
+    }
+    
+    fn parameters(&self) -> Vec<&Tensor> {
+        vec![&self.q_proj, &self.k_proj, &self.v_proj, &self.o_proj]
+    }
+    
+    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
+        vec![&mut self.q_proj, &mut self.k_proj, &mut self.v_proj, &mut self.o_proj]
+    }
+}
+
+impl Layer for LinearAttentionLayer {
+    fn forward(&self, input: &Tensor) -> Result<Tensor> {
+        self.forward(input)
+    }
+    
+    fn layer_type(&self) -> &'static str {
+        "LinearAttentionLayer"
+    }
+    
+    fn device(&self) -> &Device {
+        &self.device
+    }
+    
+    fn parameters(&self) -> Vec<&Tensor> {
+        let mut params = self.attention.parameters();
+        params.extend(self.norm.parameters());
+        params
+    }
+    
+    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
+        let mut params = self.attention.parameters_mut();
+        params.extend(self.norm.parameters_mut());
+        params
+    }
+}