355 lines
14 KiB
Rust
355 lines
14 KiB
Rust
use crate::Graph;
|
||
use crate::error::{GeomError, Result};
|
||
use crate::layers::GNNLayer;
|
||
use rtx_tensor::Tensor;
|
||
use std::collections::HashMap;
|
||
|
||
/// Graph Attention Network (GAT) Layer
|
||
///
|
||
/// Implements multi-head attention mechanism from Veličković et al. (2018):
|
||
/// α_ij = softmax(LeakyReLU(a^T [W h_i || W h_j]))
|
||
/// h_i' = σ(Σ_j α_ij W h_j)
|
||
#[derive(Debug, Clone)]
|
||
pub struct GATLayer {
|
||
weight: Tensor,
|
||
attention_weight: Tensor,
|
||
bias: Tensor,
|
||
input_dim: usize,
|
||
output_dim: usize,
|
||
num_heads: usize,
|
||
gradients_enabled: bool,
|
||
computed_gradients: bool,
|
||
attention_weights: HashMap<(usize, usize), f32>, // (from, to) -> attention weight
|
||
}
|
||
|
||
impl GATLayer {
|
||
pub fn new(input_dim: usize, output_dim: usize, num_heads: usize) -> Result<Self> {
|
||
if input_dim == 0 || output_dim == 0 || num_heads == 0 {
|
||
return Err(GeomError::LayerInitializationFailed {
|
||
reason: "Dimensions and num_heads must be greater than 0".to_string(),
|
||
});
|
||
}
|
||
|
||
if !output_dim.is_multiple_of(num_heads) {
|
||
return Err(GeomError::LayerInitializationFailed {
|
||
reason: "Output dimension must be divisible by number of heads".to_string(),
|
||
});
|
||
}
|
||
|
||
let head_dim = output_dim / num_heads;
|
||
|
||
// Xavier initialization
|
||
let fan_in = input_dim as f32;
|
||
let fan_out = output_dim as f32;
|
||
let limit = (6.0 / (fan_in + fan_out)).sqrt();
|
||
|
||
// Weight matrix for feature transformation
|
||
let mut weight_data = Vec::with_capacity(input_dim * output_dim);
|
||
for i in 0..(input_dim * output_dim) {
|
||
let val = (((i * 7919 + 1337) % 10000) as f32 / 5000.0 - 1.0) * limit;
|
||
weight_data.push(val);
|
||
}
|
||
let weight = Tensor::from_data(
|
||
weight_data,
|
||
[input_dim, output_dim],
|
||
&rtx_tensor::Device::default(),
|
||
)
|
||
.unwrap();
|
||
|
||
// Attention weight vector (2 * head_dim for concatenated features)
|
||
let attention_dim = 2 * head_dim * num_heads;
|
||
let mut attention_data = Vec::with_capacity(attention_dim);
|
||
for i in 0..attention_dim {
|
||
let val = (((i * 6151 + 2357) % 10000) as f32 / 5000.0 - 1.0) * limit;
|
||
attention_data.push(val);
|
||
}
|
||
let attention_weight = Tensor::from_data(
|
||
attention_data,
|
||
[attention_dim],
|
||
&rtx_tensor::Device::default(),
|
||
)
|
||
.unwrap();
|
||
|
||
// Bias
|
||
let bias = Tensor::zeros([output_dim], &rtx_tensor::Device::default()).unwrap();
|
||
|
||
Ok(Self {
|
||
weight,
|
||
attention_weight,
|
||
bias,
|
||
input_dim,
|
||
output_dim,
|
||
num_heads,
|
||
gradients_enabled: false,
|
||
computed_gradients: false,
|
||
attention_weights: HashMap::new(),
|
||
})
|
||
}
|
||
|
||
pub fn get_attention_weights(&self) -> &HashMap<(usize, usize), f32> {
|
||
&self.attention_weights
|
||
}
|
||
|
||
fn compute_attention(&mut self, graph: &Graph) -> Result<HashMap<usize, Vec<(usize, f32)>>> {
|
||
let mut attention_map = HashMap::new();
|
||
|
||
for node_id in graph.nodes() {
|
||
let neighbors = graph.incoming_neighbors(node_id);
|
||
let mut node_attentions = Vec::new();
|
||
|
||
if !neighbors.is_empty() {
|
||
let node_features = graph.node(node_id).unwrap().features();
|
||
|
||
let _attention_scores: Vec<f32> = Vec::new();
|
||
let mut raw_scores = Vec::new();
|
||
|
||
// Compute raw attention scores for all neighbors
|
||
for &neighbor_id in &neighbors {
|
||
let neighbor_features = graph.node(neighbor_id).unwrap().features();
|
||
|
||
// Compute attention score (simplified)
|
||
let node_2d = node_features.reshape([1, self.input_dim]).map_err(|e| {
|
||
GeomError::ForwardPassFailed {
|
||
reason: format!("Node feature reshape failed: {e}"),
|
||
}
|
||
})?;
|
||
|
||
let neighbor_2d =
|
||
neighbor_features
|
||
.reshape([1, self.input_dim])
|
||
.map_err(|e| GeomError::ForwardPassFailed {
|
||
reason: format!("Neighbor feature reshape failed: {e}"),
|
||
})?;
|
||
|
||
let node_transformed =
|
||
node_2d
|
||
.matmul(&self.weight)
|
||
.map_err(|e| GeomError::ForwardPassFailed {
|
||
reason: format!("Node transformation failed: {e}"),
|
||
})?;
|
||
|
||
let neighbor_transformed = neighbor_2d.matmul(&self.weight).map_err(|e| {
|
||
GeomError::ForwardPassFailed {
|
||
reason: format!("Neighbor transformation failed: {e}"),
|
||
}
|
||
})?;
|
||
|
||
// Simplified attention computation: dot product of first elements
|
||
let node_data = node_transformed.to_cpu().unwrap();
|
||
let neighbor_data = neighbor_transformed.to_cpu().unwrap();
|
||
|
||
let raw_score = if !node_data.is_empty() && !neighbor_data.is_empty() {
|
||
node_data[0] * neighbor_data[0] // Simplified attention mechanism
|
||
} else {
|
||
0.0
|
||
};
|
||
|
||
raw_scores.push(raw_score);
|
||
}
|
||
|
||
// Apply softmax to get attention weights
|
||
let max_score = raw_scores.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
||
let exp_scores: Vec<f32> =
|
||
raw_scores.iter().map(|&s| (s - max_score).exp()).collect();
|
||
let sum_exp = exp_scores.iter().sum::<f32>();
|
||
|
||
for (i, &neighbor_id) in neighbors.iter().enumerate() {
|
||
let attention_weight = if sum_exp > 0.0 {
|
||
exp_scores[i] / sum_exp
|
||
} else {
|
||
1.0 / neighbors.len() as f32
|
||
};
|
||
|
||
node_attentions.push((neighbor_id.index(), attention_weight));
|
||
|
||
// Store attention weights for inspection
|
||
self.attention_weights
|
||
.insert((node_id.index(), neighbor_id.index()), attention_weight);
|
||
}
|
||
}
|
||
|
||
attention_map.insert(node_id.index(), node_attentions);
|
||
}
|
||
|
||
Ok(attention_map)
|
||
}
|
||
}
|
||
|
||
impl GNNLayer for GATLayer {
|
||
fn forward(&mut self, graph: &Graph) -> Result<Vec<Tensor>> {
|
||
if graph.node_count() == 0 {
|
||
return Err(GeomError::ForwardPassFailed {
|
||
reason: "Cannot process empty graph".to_string(),
|
||
});
|
||
}
|
||
|
||
// Compute attention weights
|
||
let attention_map = self.compute_attention(graph)?;
|
||
|
||
let mut outputs = Vec::with_capacity(graph.node_count());
|
||
|
||
for node_id in graph.nodes() {
|
||
let node_idx = node_id.index();
|
||
|
||
if let Some(attentions) = attention_map.get(&node_idx) {
|
||
if attentions.is_empty() {
|
||
// Isolated node - transform its own features
|
||
let node_features = graph.node(node_id).unwrap().features();
|
||
let node_2d = node_features.reshape([1, self.input_dim]).map_err(|e| {
|
||
GeomError::ForwardPassFailed {
|
||
reason: format!("Self feature reshape failed: {e}"),
|
||
}
|
||
})?;
|
||
|
||
let transformed =
|
||
node_2d
|
||
.matmul(&self.weight)
|
||
.map_err(|e| GeomError::ForwardPassFailed {
|
||
reason: format!("Self transformation failed: {e}"),
|
||
})?;
|
||
let biased =
|
||
transformed
|
||
.add(&self.bias)
|
||
.map_err(|e| GeomError::ForwardPassFailed {
|
||
reason: format!("Bias addition failed: {e}"),
|
||
})?;
|
||
|
||
// Reshape back to [output_dim] from [1, output_dim]
|
||
let output = biased.reshape([self.output_dim]).map_err(|e| {
|
||
GeomError::ForwardPassFailed {
|
||
reason: format!("Output reshape failed: {e}"),
|
||
}
|
||
})?;
|
||
|
||
outputs.push(output);
|
||
} else {
|
||
// Aggregate neighbor features with attention weights
|
||
let mut aggregated = None;
|
||
|
||
for &(neighbor_idx, attention_weight) in attentions {
|
||
// Find neighbor node by index
|
||
let neighbor_id = graph
|
||
.nodes()
|
||
.into_iter()
|
||
.find(|id| id.index() == neighbor_idx)
|
||
.ok_or_else(|| GeomError::ForwardPassFailed {
|
||
reason: format!("Neighbor node {neighbor_idx} not found"),
|
||
})?;
|
||
|
||
let neighbor_features = graph.node(neighbor_id).unwrap().features();
|
||
let neighbor_2d =
|
||
neighbor_features
|
||
.reshape([1, self.input_dim])
|
||
.map_err(|e| GeomError::ForwardPassFailed {
|
||
reason: format!("Neighbor feature reshape failed: {e}"),
|
||
})?;
|
||
|
||
let transformed = neighbor_2d.matmul(&self.weight).map_err(|e| {
|
||
GeomError::ForwardPassFailed {
|
||
reason: format!("Neighbor transformation failed: {e}"),
|
||
}
|
||
})?;
|
||
|
||
let weighted = transformed
|
||
.mul(
|
||
&Tensor::from_data(
|
||
vec![attention_weight],
|
||
[1],
|
||
&rtx_tensor::Device::default(),
|
||
)
|
||
.unwrap(),
|
||
)
|
||
.map_err(|e| GeomError::ForwardPassFailed {
|
||
reason: format!("Attention weighting failed: {e}"),
|
||
})?;
|
||
|
||
aggregated = match aggregated {
|
||
None => Some(weighted),
|
||
Some(acc) => Some(acc.add(&weighted).map_err(|e| {
|
||
GeomError::ForwardPassFailed {
|
||
reason: format!("Aggregation failed: {e}"),
|
||
}
|
||
})?),
|
||
};
|
||
}
|
||
|
||
let output = if let Some(agg) = aggregated {
|
||
agg.add(&self.bias)
|
||
.map_err(|e| GeomError::ForwardPassFailed {
|
||
reason: format!("Bias addition failed: {e}"),
|
||
})?
|
||
} else {
|
||
return Err(GeomError::ForwardPassFailed {
|
||
reason: "No aggregated features computed".to_string(),
|
||
});
|
||
};
|
||
|
||
// Apply ELU activation (approximated with ReLU for simplicity)
|
||
let activated = output.relu().map_err(|e| GeomError::ForwardPassFailed {
|
||
reason: format!("Activation failed: {e}"),
|
||
})?;
|
||
|
||
// Reshape back to [output_dim] from [1, output_dim]
|
||
let output_1d = activated.reshape([self.output_dim]).map_err(|e| {
|
||
GeomError::ForwardPassFailed {
|
||
reason: format!("Output reshape failed: {e}"),
|
||
}
|
||
})?;
|
||
|
||
outputs.push(output_1d);
|
||
}
|
||
} else {
|
||
return Err(GeomError::ForwardPassFailed {
|
||
reason: format!("No attention computed for node {node_idx}"),
|
||
});
|
||
}
|
||
}
|
||
|
||
Ok(outputs)
|
||
}
|
||
|
||
fn parameters(&self) -> Vec<Tensor> {
|
||
vec![
|
||
self.weight.clone(),
|
||
self.attention_weight.clone(),
|
||
self.bias.clone(),
|
||
]
|
||
}
|
||
|
||
fn enable_gradients(&mut self, enabled: bool) {
|
||
self.gradients_enabled = enabled;
|
||
}
|
||
|
||
fn has_computed_gradients(&self) -> bool {
|
||
self.computed_gradients
|
||
}
|
||
|
||
fn backward(&mut self, grad_outputs: &[Tensor]) -> Result<Vec<Tensor>> {
|
||
if !self.gradients_enabled {
|
||
return Err(GeomError::BackwardPassFailed {
|
||
reason: "Gradients are not enabled for this layer".to_string(),
|
||
});
|
||
}
|
||
|
||
if grad_outputs.is_empty() {
|
||
return Err(GeomError::BackwardPassFailed {
|
||
reason: "No gradient outputs provided".to_string(),
|
||
});
|
||
}
|
||
|
||
// Simplified gradient computation
|
||
let weight_grad =
|
||
Tensor::zeros(self.weight.shape().dims(), &rtx_tensor::Device::default()).unwrap();
|
||
let attention_grad = Tensor::zeros(
|
||
[self.attention_weight.shape().dims()[0]],
|
||
&rtx_tensor::Device::default(),
|
||
)
|
||
.unwrap();
|
||
let bias_grad = Tensor::zeros([self.output_dim], &rtx_tensor::Device::default()).unwrap();
|
||
|
||
self.computed_gradients = true;
|
||
|
||
Ok(vec![weight_grad, attention_grad, bias_grad])
|
||
}
|
||
}
|