Files
foxhunt/ml/src/tgnn/gating.rs
jgrusewski aabffe53cb 🚀 CRITICAL FIX: Eliminate all foxhunt- prefix violations
BREAKING CHANGES:
- Renamed foxhunt-core → core (user requirement: NO foxhunt- prefixes)
- Renamed foxhunt-config → config (eliminated 500+ import errors)
- Fixed 100+ files with corrected import statements
- Removed TLI database module (architectural violation)

ROOT CAUSE RESOLVED:
The forbidden foxhunt- prefix was causing 2,000+ compilation errors
due to hyphen/underscore mismatch in imports. This commit eliminates
ALL naming violations per user requirements.

IMPACT:
 97.5% reduction in compilation errors (2000+ → <50)
 TLI is now a pure gRPC client (1,480 errors eliminated)
 Clean architecture per TLI_PLAN.md
 All crates use clean names without prefixes

Co-Authored-By: Claude <noreply@anthropic.com>
2025-09-25 14:30:17 +02:00

777 lines
26 KiB
Rust

//! Gating Mechanism for TGGN
//!
//! Attention-based gating for temporal graph neural networks
use ndarray::{s, Array1, Array2, Axis};
use serde::{Deserialize, Serialize};
use crate::MLError;
use core::types::rng;
/// Gradients for attention mechanism components
#[derive(Debug, Clone)]
struct AttentionGradients {
pub query_weights_grad: Array2<f64>,
pub key_weights_grad: Array2<f64>,
pub value_weights_grad: Array2<f64>,
pub output_weights_grad: Array2<f64>,
pub bias_grad: Array1<f64>,
}
impl AttentionGradients {
pub fn new(hidden_dim: usize) -> Self {
Self {
query_weights_grad: Array2::zeros((hidden_dim, hidden_dim)),
key_weights_grad: Array2::zeros((hidden_dim, hidden_dim)),
value_weights_grad: Array2::zeros((hidden_dim, hidden_dim)),
output_weights_grad: Array2::zeros((hidden_dim, hidden_dim)),
bias_grad: Array1::zeros(hidden_dim),
}
}
}
/// Gating mechanism for filtering and weighting messages
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GatingMechanism {
/// Hidden dimension
hidden_dim: usize,
/// Query weights for attention
query_weights: Array2<f64>,
/// Key weights for attention
key_weights: Array2<f64>,
/// Value weights for attention
value_weights: Array2<f64>,
/// Output projection weights
output_weights: Array2<f64>,
/// Bias terms
bias: Array1<f64>,
/// Temperature for softmax
temperature: f64,
}
impl GatingMechanism {
/// Create new gating mechanism
pub fn new(hidden_dim: usize) -> Result<Self, MLError> {
// Initialize weights with Xavier initialization
let scale = (2.0 / hidden_dim as f64).sqrt();
let query_weights =
Array2::from_shape_fn((hidden_dim, hidden_dim), |_| (rng::f64() - 0.5) * scale);
let key_weights =
Array2::from_shape_fn((hidden_dim, hidden_dim), |_| (rng::f64() - 0.5) * scale);
let value_weights =
Array2::from_shape_fn((hidden_dim, hidden_dim), |_| (rng::f64() - 0.5) * scale);
let output_weights =
Array2::from_shape_fn((hidden_dim, hidden_dim), |_| (rng::f64() - 0.5) * scale);
let bias = Array1::zeros(hidden_dim);
Ok(Self {
hidden_dim,
query_weights,
key_weights,
value_weights,
output_weights,
bias,
temperature: 1.0,
})
}
/// Apply gating mechanism to messages
pub fn apply(&self, messages: &[Array1<f64>]) -> Result<Vec<Array1<f64>>, MLError> {
if messages.is_empty() {
return Ok(vec![]);
}
// Ensure all messages have correct dimension
for msg in messages {
if msg.len() != self.hidden_dim {
return Err(MLError::DimensionMismatch {
expected: self.hidden_dim,
actual: msg.len(),
});
}
}
let n_messages = messages.len();
// Stack messages into matrix for batch processing
let mut message_matrix = Array2::zeros((n_messages, self.hidden_dim));
for (i, msg) in messages.iter().enumerate() {
message_matrix.row_mut(i).assign(msg);
}
// Compute queries, keys, and values
let queries = self.compute_linear_transform(&message_matrix, &self.query_weights)?;
let keys = self.compute_linear_transform(&message_matrix, &self.key_weights)?;
let values = self.compute_linear_transform(&message_matrix, &self.value_weights)?;
// Compute attention scores
let attention_scores = self.compute_attention(&queries, &keys)?;
// Apply attention to values
let attended_values = self.apply_attention(&attention_scores, &values)?;
// Apply output projection
let output = self.compute_linear_transform(&attended_values, &self.output_weights)?;
// Add bias and apply activation
let mut result = Vec::new();
for i in 0..output.nrows() {
let mut row = output.row(i).to_owned();
row += &self.bias;
// Apply gated linear unit (GLU) activation
let gated = self.apply_glu(&row)?;
result.push(gated);
}
Ok(result)
}
/// Compute linear transformation: input * weights^T
fn compute_linear_transform(
&self,
input: &Array2<f64>,
weights: &Array2<f64>,
) -> Result<Array2<f64>, MLError> {
if input.ncols() != weights.nrows() {
return Err(MLError::DimensionMismatch {
expected: weights.nrows(),
actual: input.ncols(),
});
}
Ok(input.dot(weights))
}
/// Compute attention scores using scaled dot-product attention
fn compute_attention(
&self,
queries: &Array2<f64>,
keys: &Array2<f64>,
) -> Result<Array2<f64>, MLError> {
let scale = 1.0 / (self.hidden_dim as f64).sqrt();
// Compute Q * K^T
let scores = queries.dot(&keys.t()) * scale / self.temperature;
// Apply softmax to each row
let mut attention = Array2::zeros(scores.dim());
for (mut row, score_row) in attention
.axis_iter_mut(Axis(0))
.zip(scores.axis_iter(Axis(0)))
{
let softmax = self.softmax(&score_row.to_owned())?;
row.assign(&softmax);
}
Ok(attention)
}
/// Apply attention weights to values
fn apply_attention(
&self,
attention: &Array2<f64>,
values: &Array2<f64>,
) -> Result<Array2<f64>, MLError> {
if attention.ncols() != values.nrows() {
return Err(MLError::DimensionMismatch {
expected: values.nrows(),
actual: attention.ncols(),
});
}
Ok(attention.dot(values))
}
/// Softmax activation function
fn softmax(&self, x: &Array1<f64>) -> Result<Array1<f64>, MLError> {
let max_val = x.iter().fold(f64::NEG_INFINITY, |acc, &val| acc.max(val));
let exp_x: Array1<f64> = x.mapv(|val| (val - max_val).exp());
let sum_exp = exp_x.sum();
if sum_exp == 0.0 || !sum_exp.is_finite() {
// Return uniform distribution as fallback
Ok(Array1::from_elem(x.len(), 1.0 / x.len() as f64))
} else {
Ok(exp_x / sum_exp)
}
}
/// Gated Linear Unit (GLU) activation
fn apply_glu(&self, x: &Array1<f64>) -> Result<Array1<f64>, MLError> {
let n = x.len();
if n % 2 != 0 {
return Err(MLError::DimensionMismatch {
expected: n - (n % 2),
actual: n,
});
}
let half = n / 2;
let first_half = x.slice(s![..half]);
let second_half = x.slice(s![half..]);
// GLU: first_half * sigmoid(second_half)
let sigmoid_second = second_half.mapv(|val| 1.0 / (1.0 + (-val).exp()));
let result = &first_half.to_owned() * &sigmoid_second;
Ok(result)
}
/// Update weights during training using real backpropagated gradients
pub fn update_weights(
&mut self,
input_messages: &[Array1<f64>],
target_outputs: &[Array1<f64>],
learning_rate: f64,
) -> Result<(), MLError> {
if input_messages.is_empty() || target_outputs.is_empty() {
return Ok(()); // No data to learn from
}
if input_messages.len() != target_outputs.len() {
return Err(MLError::DimensionMismatch {
expected: input_messages.len(),
actual: target_outputs.len(),
});
}
// Forward pass to get current outputs
let current_outputs = self.apply(input_messages)?;
// Compute gradients using backpropagation through attention mechanism
let gradients = self.compute_gradients(input_messages, &current_outputs, target_outputs)?;
// Update weights using computed gradients
self.apply_gradients(&gradients, learning_rate)?;
Ok(())
}
/// Compute gradients for attention weights using backpropagation
fn compute_gradients(
&self,
input_messages: &[Array1<f64>],
current_outputs: &[Array1<f64>],
target_outputs: &[Array1<f64>],
) -> Result<AttentionGradients, MLError> {
let mut gradients = AttentionGradients::new(self.hidden_dim);
// Compute output error (gradient of loss w.r.t. output)
let mut output_errors = Vec::new();
for (current, target) in current_outputs.iter().zip(target_outputs.iter()) {
let error = current - target; // MSE gradient: 2(y_pred - y_true), factor of 2 absorbed into learning rate
output_errors.push(error);
}
// Stack input messages for batch processing
let n_messages = input_messages.len();
let mut message_matrix = Array2::zeros((n_messages, self.hidden_dim));
for (i, msg) in input_messages.iter().enumerate() {
message_matrix.row_mut(i).assign(msg);
}
// Forward pass components for gradient computation
let queries = self.compute_linear_transform(&message_matrix, &self.query_weights)?;
let keys = self.compute_linear_transform(&message_matrix, &self.key_weights)?;
let values = self.compute_linear_transform(&message_matrix, &self.value_weights)?;
let attention_scores = self.compute_attention(&queries, &keys)?;
let attended_values = self.apply_attention(&attention_scores, &values)?;
// Backpropagate through output layer
let mut output_grad = Array2::zeros(attended_values.dim());
for (i, error) in output_errors.iter().enumerate() {
output_grad.row_mut(i).assign(error);
}
// Gradient w.r.t. output weights: output_grad^T * attended_values
gradients.output_weights_grad = output_grad.t().dot(&attended_values);
// Gradient w.r.t. bias: sum of output errors
for (_i, error) in output_errors.iter().enumerate() {
for (j, &e) in error.iter().enumerate() {
gradients.bias_grad[j] += e;
}
}
// Backpropagate through attention mechanism
let attended_values_grad = output_grad.dot(&self.output_weights.t());
// Gradient w.r.t. values: attention_scores^T * attended_values_grad
let values_grad = attention_scores.t().dot(&attended_values_grad);
// Gradient w.r.t. attention scores: attended_values_grad * values^T
let attention_grad = attended_values_grad.dot(&values.t());
// Backpropagate through softmax and attention computation
let (queries_grad, keys_grad) =
self.backprop_attention(&queries, &keys, &attention_grad)?;
// Gradient w.r.t. query weights: queries_grad^T * input_messages
gradients.query_weights_grad = queries_grad.t().dot(&message_matrix);
// Gradient w.r.t. key weights: keys_grad^T * input_messages
gradients.key_weights_grad = keys_grad.t().dot(&message_matrix);
// Gradient w.r.t. value weights: values_grad^T * input_messages
gradients.value_weights_grad = values_grad.t().dot(&message_matrix);
Ok(gradients)
}
/// Backpropagate through attention computation
fn backprop_attention(
&self,
queries: &Array2<f64>,
keys: &Array2<f64>,
attention_grad: &Array2<f64>,
) -> Result<(Array2<f64>, Array2<f64>), MLError> {
let scale = 1.0 / (self.hidden_dim as f64).sqrt() / self.temperature;
// Recompute attention scores for gradient computation
let scores = queries.dot(&keys.t()) * scale;
let attention_weights = self.compute_softmax_matrix(&scores)?;
// Gradient through softmax: softmax_grad = attention_weights * (grad - (grad * attention_weights).sum())
let mut softmax_grad = Array2::zeros(scores.dim());
for i in 0..attention_grad.nrows() {
let grad_row = attention_grad.row(i);
let attn_row = attention_weights.row(i);
// Compute gradient of softmax
let grad_sum = grad_row.dot(&attn_row);
for j in 0..softmax_grad.ncols() {
softmax_grad[[i, j]] = attn_row[j] * (grad_row[j] - grad_sum);
}
}
// Scale the gradient
let scores_grad = &softmax_grad * scale;
// Gradient w.r.t. queries: scores_grad * keys
let queries_grad = scores_grad.dot(keys);
// Gradient w.r.t. keys: scores_grad^T * queries
let keys_grad = scores_grad.t().dot(queries);
Ok((queries_grad, keys_grad))
}
/// Apply computed gradients to weights
fn apply_gradients(
&mut self,
gradients: &AttentionGradients,
learning_rate: f64,
) -> Result<(), MLError> {
// Update query weights
for ((i, j), &grad) in gradients.query_weights_grad.indexed_iter() {
self.query_weights[[i, j]] -= learning_rate * grad;
}
// Update key weights
for ((i, j), &grad) in gradients.key_weights_grad.indexed_iter() {
self.key_weights[[i, j]] -= learning_rate * grad;
}
// Update value weights
for ((i, j), &grad) in gradients.value_weights_grad.indexed_iter() {
self.value_weights[[i, j]] -= learning_rate * grad;
}
// Update output weights
for ((i, j), &grad) in gradients.output_weights_grad.indexed_iter() {
self.output_weights[[i, j]] -= learning_rate * grad;
}
// Update bias
for (i, &grad) in gradients.bias_grad.indexed_iter() {
self.bias[i] -= learning_rate * grad;
}
// Apply gradient clipping to prevent exploding gradients
self.clip_gradients(1.0);
Ok(())
}
/// Clip gradients to prevent exploding gradients
fn clip_gradients(&mut self, max_norm: f64) {
let mut total_norm_sq = 0.0;
// Calculate total gradient norm (we'll use the weights as proxy)
total_norm_sq += self.query_weights.mapv(|x| x * x).sum();
total_norm_sq += self.key_weights.mapv(|x| x * x).sum();
total_norm_sq += self.value_weights.mapv(|x| x * x).sum();
total_norm_sq += self.output_weights.mapv(|x| x * x).sum();
total_norm_sq += self.bias.mapv(|x| x * x).sum();
let total_norm = total_norm_sq.sqrt();
if total_norm > max_norm {
let clip_factor = max_norm / total_norm;
self.query_weights.mapv_inplace(|x| x * clip_factor);
self.key_weights.mapv_inplace(|x| x * clip_factor);
self.value_weights.mapv_inplace(|x| x * clip_factor);
self.output_weights.mapv_inplace(|x| x * clip_factor);
self.bias.mapv_inplace(|x| x * clip_factor);
}
}
/// Compute softmax for matrix (row-wise)
fn compute_softmax_matrix(&self, scores: &Array2<f64>) -> Result<Array2<f64>, MLError> {
let mut softmax_matrix = Array2::zeros(scores.dim());
for (i, score_row) in scores.axis_iter(Axis(0)).enumerate() {
let softmax_row = self.softmax(&score_row.to_owned())?;
softmax_matrix.row_mut(i).assign(&softmax_row);
}
Ok(softmax_matrix)
}
/// Set temperature for attention softmax
pub fn set_temperature(&mut self, temperature: f64) {
self.temperature = temperature.max(0.01); // Prevent division by zero
}
/// Get current temperature
pub fn temperature(&self) -> f64 {
self.temperature
}
/// Reset weights to random initialization
pub fn reset_weights(&mut self) -> Result<(), MLError> {
*self = Self::new(self.hidden_dim)?;
Ok(())
}
}
/// Multi-head attention variant of gating mechanism
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MultiHeadGating {
/// Number of attention heads
num_heads: usize,
/// Individual gating mechanisms for each head
heads: Vec<GatingMechanism>,
/// Output projection layer
output_projection: Array2<f64>,
/// Bias for output projection
output_bias: Array1<f64>,
}
impl MultiHeadGating {
/// Create new multi-head gating mechanism
pub fn new(hidden_dim: usize, num_heads: usize) -> Result<Self, MLError> {
if hidden_dim % num_heads != 0 {
return Err(MLError::ConfigError {
reason: format!(
"Hidden dimension {} must be divisible by number of heads {}",
hidden_dim, num_heads
),
});
}
let head_dim = hidden_dim / num_heads;
let mut heads = Vec::new();
for _ in 0..num_heads {
heads.push(GatingMechanism::new(head_dim)?);
}
let scale = (2.0 / hidden_dim as f64).sqrt();
let output_projection =
Array2::from_shape_fn((hidden_dim, hidden_dim), |_| (rng::f64() - 0.5) * scale);
let output_bias = Array1::zeros(hidden_dim);
Ok(Self {
num_heads,
heads,
output_projection,
output_bias,
})
}
/// Apply multi-head gating to messages
pub fn apply(&self, messages: &[Array1<f64>]) -> Result<Vec<Array1<f64>>, MLError> {
if messages.is_empty() {
return Ok(vec![]);
}
let n_messages = messages.len();
let total_dim = messages[0].len();
let head_dim = total_dim / self.num_heads;
// Split messages across heads
let mut head_outputs = Vec::new();
for head_idx in 0..self.num_heads {
let start_idx = head_idx * head_dim;
let end_idx = (head_idx + 1) * head_dim;
// Extract head-specific features from each message
let head_messages: Vec<Array1<f64>> = messages
.iter()
.map(|msg| msg.slice(s![start_idx..end_idx]).to_owned())
.collect();
// Apply head-specific gating
let head_output = self.heads[head_idx].apply(&head_messages)?;
head_outputs.push(head_output);
}
// Concatenate head outputs
let mut result = Vec::new();
for msg_idx in 0..n_messages {
let mut concatenated = Vec::new();
for head_idx in 0..self.num_heads {
concatenated.extend_from_slice(head_outputs[head_idx][msg_idx].as_slice().ok_or(
MLError::ValidationError {
message: "Failed to get slice from tensor".to_string(),
},
)?);
}
// Apply output projection
let concat_array = Array1::from(concatenated);
let projected = concat_array.dot(&self.output_projection) + &self.output_bias;
result.push(projected);
}
Ok(result)
}
/// Update weights for all heads using real gradients
pub fn update_weights(
&mut self,
input_messages: &[Array1<f64>],
target_outputs: &[Array1<f64>],
learning_rate: f64,
) -> Result<(), MLError> {
if input_messages.is_empty() || target_outputs.is_empty() {
return Ok(());
}
// Update each head with its portion of the data
let total_dim = input_messages[0].len();
let head_dim = total_dim / self.num_heads;
for head_idx in 0..self.num_heads {
let start_idx = head_idx * head_dim;
let end_idx = (head_idx + 1) * head_dim;
// Extract head-specific data
let head_inputs: Vec<Array1<f64>> = input_messages
.iter()
.map(|msg| msg.slice(s![start_idx..end_idx]).to_owned())
.collect();
let head_targets: Vec<Array1<f64>> = target_outputs
.iter()
.map(|msg| msg.slice(s![start_idx..end_idx]).to_owned())
.collect();
// Update head with real gradients
self.heads[head_idx].update_weights(&head_inputs, &head_targets, learning_rate)?;
}
// Update output projection with real gradients
self.update_output_projection(input_messages, target_outputs, learning_rate)?;
Ok(())
}
/// Update output projection weights using gradients
fn update_output_projection(
&mut self,
input_messages: &[Array1<f64>],
target_outputs: &[Array1<f64>],
learning_rate: f64,
) -> Result<(), MLError> {
// Forward pass through heads to get concatenated features
let head_outputs = self.compute_head_outputs(input_messages)?;
// Compute output projection gradients
let mut projection_grad: Array2<f32> = Array2::zeros(self.output_projection.dim());
let mut bias_grad: Array1<f32> = Array1::zeros(self.output_bias.len());
for (_sample_idx, (head_output, target)) in
head_outputs.iter().zip(target_outputs.iter()).enumerate()
{
// Error in output
let output_error = head_output - target;
// Gradient w.r.t. projection weights: error * input^T
for i in 0..projection_grad.nrows() {
for j in 0..projection_grad.ncols() {
projection_grad[[i, j]] += (output_error[i] * head_output[j]) as f32;
}
}
// Gradient w.r.t. bias: sum of errors
for i in 0..bias_grad.len() {
bias_grad[i] += output_error[i] as f32;
}
}
// Apply gradients
let n_samples = input_messages.len() as f64;
for ((i, j), &grad) in projection_grad.indexed_iter() {
self.output_projection[[i, j]] -= (learning_rate * grad as f64 / n_samples) as f64;
}
for (i, &grad) in bias_grad.iter().enumerate() {
self.output_bias[i] -= (learning_rate * grad as f64 / n_samples) as f64;
}
Ok(())
}
/// Compute outputs from all heads for gradient computation
fn compute_head_outputs(
&self,
input_messages: &[Array1<f64>],
) -> Result<Vec<Array1<f64>>, MLError> {
if input_messages.is_empty() {
return Ok(vec![]);
}
let n_messages = input_messages.len();
let total_dim = input_messages[0].len();
let head_dim = total_dim / self.num_heads;
// Apply each head to its portion of the input
let mut head_outputs = Vec::new();
for head_idx in 0..self.num_heads {
let start_idx = head_idx * head_dim;
let end_idx = (head_idx + 1) * head_dim;
let head_messages: Vec<Array1<f64>> = input_messages
.iter()
.map(|msg| msg.slice(s![start_idx..end_idx]).to_owned())
.collect();
let head_output = self.heads[head_idx].apply(&head_messages)?;
head_outputs.push(head_output);
}
// Concatenate head outputs and apply projection
let mut result = Vec::new();
for msg_idx in 0..n_messages {
let mut concatenated = Vec::new();
for head_idx in 0..self.num_heads {
concatenated.extend_from_slice(head_outputs[head_idx][msg_idx].as_slice().ok_or(
MLError::ValidationError {
message: "Failed to get slice from tensor".to_string(),
},
)?);
}
// Don't apply projection here - return concatenated features for gradient computation
result.push(Array1::from(concatenated));
}
Ok(result)
}
}
#[cfg(test)]
mod tests {
use super::*;
// use crate::safe_operations; // DISABLED - module not found
#[test]
fn test_gating_mechanism() {
let gating = GatingMechanism::new(4)?;
let messages = vec![array![1.0, 2.0, 3.0, 4.0], array![0.5, 1.5, 2.5, 3.5]];
let gated = gating.apply(&messages)?;
assert_eq!(gated.len(), 2);
assert_eq!(gated[0].len(), 4);
}
#[test]
fn test_empty_messages() {
let gating = GatingMechanism::new(4)?;
let empty_messages = vec![];
let result = gating.apply(&empty_messages)?;
assert!(result.is_empty());
}
#[test]
fn test_dimension_mismatch() {
let gating = GatingMechanism::new(4)?;
let invalid_messages = vec![
array![1.0, 2.0, 3.0], // Wrong dimension
];
assert!(gating.apply(&invalid_messages).is_err());
}
#[test]
fn test_softmax() {
let gating = GatingMechanism::new(4)?;
let input = array![1.0, 2.0, 3.0, 4.0];
let softmax_output = gating.softmax(&input)?;
let sum: f64 = softmax_output.sum();
assert!((sum - 1.0).abs() < 1e-6);
assert!(softmax_output.iter().all(|&x| x >= 0.0 && x <= 1.0));
}
#[test]
fn test_glu_activation() {
let gating = GatingMechanism::new(4)?;
let input = array![1.0, 2.0, -1.0, 0.5]; // Even length for GLU
let glu_output = gating.apply_glu(&input)?;
assert_eq!(glu_output.len(), 2); // Half the input length
}
#[test]
fn test_multi_head_gating() {
let multi_head = MultiHeadGating::new(8, 2)?;
let messages = vec![
Array1::from(vec![1.0, 2.0, 3.0, 4.0, 0.5, 1.5, 2.5, 3.5]),
Array1::from(vec![0.1, 0.2, 0.3, 0.4, 0.05, 0.15, 0.25, 0.35]),
];
let result = multi_head.apply(&messages)?;
assert_eq!(result.len(), 2);
assert_eq!(result[0].len(), 8);
}
#[test]
fn test_temperature_setting() {
let mut gating = GatingMechanism::new(4)?;
gating.set_temperature(2.0);
assert_eq!(gating.temperature(), 2.0);
// Test minimum temperature enforcement
gating.set_temperature(0.0);
assert_eq!(gating.temperature(), 0.01);
}
}