Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,219 @@
|
||||
//! KAN Networks (Kolmogorov-Arnold Networks) Implementation
|
||||
//!
|
||||
//! This module implements KAN networks based on the Kolmogorov-Arnold representation theorem.
|
||||
//! Unlike traditional MLPs that use fixed activations on nodes, KANs use learnable activation
|
||||
//! functions on edges, enabling superior function approximation and interpretability.
|
||||
//!
|
||||
//! ## Key Features
|
||||
//!
|
||||
//! - **Edge-based learnable activations**: Each edge has its own trainable activation function
|
||||
//! - **Multiple basis functions**: Support for B-splines, Chebyshev polynomials, and Fourier series
|
||||
//! - **Adaptive grid refinement**: Automatic grid extension based on function complexity
|
||||
//! - **Network pruning**: Interpretability-guided pruning for simplified networks
|
||||
//! - **Symbolic regression**: Extract symbolic mathematical expressions from trained networks
|
||||
//! - **Autograd integration**: Full gradient support for training
|
||||
//!
|
||||
//! ## Architecture Overview
|
||||
//!
|
||||
//! ```text
|
||||
//! Input Layer Hidden Layer Output Layer
|
||||
//! x1 ────────────○─────────────── y1
|
||||
//! │ φ(1,1) / \ φ(2,1) │
|
||||
//! └────────────○───○──────────── y2
|
||||
//! x2 ────────────○───○────────────
|
||||
//! φ(1,2) │ │ φ(2,2)
|
||||
//! ... ...
|
||||
//! ```
|
||||
//!
|
||||
//! Where each φ(i,j) represents a learnable univariate activation function on the edge
|
||||
//! from node i in layer L to node j in layer L+1.
|
||||
//!
|
||||
//! ## Usage Example
|
||||
//!
|
||||
//! ```rust
|
||||
//! use rtx_transformers::kan::*;
|
||||
//! use rtx_tensor::{Tensor, Device};
|
||||
//!
|
||||
//! let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
//!
|
||||
//! // Create a KAN network with architecture [2, 5, 3, 1]
|
||||
//! let kan_network = KANNetwork::new(
|
||||
//! &[2, 5, 3, 1], // architecture
|
||||
//! 20, // grid_size
|
||||
//! 3, // spline_order (for B-splines)
|
||||
//! BasisFunctionType::BSpline,
|
||||
//! &device
|
||||
//! )?;
|
||||
//!
|
||||
//! // Forward pass
|
||||
//! let input = Tensor::randn(&[batch_size, 2], &device)?;
|
||||
//! let output = kan_network.forward(&input)?;
|
||||
//!
|
||||
//! // Train the network (with your training loop)
|
||||
//! // ...
|
||||
//!
|
||||
//! // Extract symbolic representation
|
||||
//! let regressor = SymbolicRegressor::new(
|
||||
//! 5, // max_depth
|
||||
//! vec![SymbolicOp::Add, SymbolicOp::Mul, SymbolicOp::Sin],
|
||||
//! 0.01, // tolerance
|
||||
//! 1000, // max_iterations
|
||||
//! &device
|
||||
//! )?;
|
||||
//! let symbolic_form = regressor.extract_from_kan(&kan_network, &train_data)?;
|
||||
//! ```
|
||||
|
||||
pub mod kan_network;
|
||||
pub mod basis_functions;
|
||||
pub mod grid_adapter;
|
||||
pub mod pruning;
|
||||
pub mod symbolic;
|
||||
|
||||
#[cfg(all(test, feature = "disabled_tests"))]
|
||||
pub mod kan_network_tests;
|
||||
|
||||
// Re-export main types for convenience
|
||||
pub use kan_network::{KANNetwork, KANLayer, KANConfig, KANState, KANGates};
|
||||
pub use basis_functions::{BasisFunction, BasisFunctionType};
|
||||
pub use grid_adapter::{GridAdapter, AdaptationStrategy};
|
||||
pub use pruning::{KANPruner, PruningStrategy};
|
||||
pub use symbolic::{SymbolicRegressor, LawDiscoverer, SymbolicOp};
|
||||
|
||||
use crate::error::{TransformerError, Result};
|
||||
|
||||
/// KAN-specific error types
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum KANError {
|
||||
#[error("Invalid grid configuration: {msg}")]
|
||||
InvalidGrid { msg: String },
|
||||
|
||||
#[error("Basis function error: {msg}")]
|
||||
BasisFunction { msg: String },
|
||||
|
||||
#[error("Grid adaptation failed: {msg}")]
|
||||
GridAdaptation { msg: String },
|
||||
|
||||
#[error("Pruning operation failed: {msg}")]
|
||||
Pruning { msg: String },
|
||||
|
||||
#[error("Symbolic regression failed: {msg}")]
|
||||
SymbolicRegression { msg: String },
|
||||
|
||||
#[error("Dimension mismatch: expected {expected}, got {actual}")]
|
||||
DimensionMismatch { expected: usize, actual: usize },
|
||||
|
||||
#[error("Autograd integration error: {msg}")]
|
||||
AutogradIntegration { msg: String },
|
||||
}
|
||||
|
||||
impl From<KANError> for TransformerError {
|
||||
fn from(err: KANError) -> Self {
|
||||
TransformerError::KAN(err.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Utility functions for KAN networks
|
||||
pub mod utils {
|
||||
use super::*;
|
||||
use rtx_tensor::{Tensor, Device};
|
||||
|
||||
/// Generate a grid of points for basis function evaluation
|
||||
pub fn generate_grid(min: f64, max: f64, size: usize, device: &Device) -> Result<Tensor> {
|
||||
if size < 2 {
|
||||
return Err(TransformerError::KAN(
|
||||
"Grid size must be at least 2".to_string()
|
||||
));
|
||||
}
|
||||
|
||||
Tensor::linspace(min, max, size, device)
|
||||
.map_err(|e| TransformerError::TensorError(e))
|
||||
}
|
||||
|
||||
/// Compute complexity measure for grid adaptation
|
||||
pub fn compute_complexity_measure(values: &Tensor) -> Result<f64> {
|
||||
// Use second derivative as complexity measure
|
||||
if values.dim() < 1 {
|
||||
return Ok(0.0);
|
||||
}
|
||||
|
||||
let n = values.shape().dims()[values.dim() - 1];
|
||||
if n < 3 {
|
||||
return Ok(0.0);
|
||||
}
|
||||
|
||||
// Approximate second derivative using finite differences
|
||||
let second_deriv = values.narrow(values.dim() - 1, 2, n - 2)?
|
||||
.sub(&values.narrow(values.dim() - 1, 1, n - 2)?.mul_scalar(2.0)?)?
|
||||
.add(&values.narrow(values.dim() - 1, 0, n - 2)?)?;
|
||||
|
||||
let complexity = second_deriv.abs()?.mean(None, false)?;
|
||||
Ok(complexity.get_item([])?)
|
||||
}
|
||||
|
||||
/// Check if a function is approximately linear
|
||||
pub fn is_approximately_linear(values: &Tensor, tolerance: f64) -> Result<bool> {
|
||||
let complexity = compute_complexity_measure(values)?;
|
||||
Ok(complexity < tolerance)
|
||||
}
|
||||
|
||||
/// Estimate optimal grid size based on function complexity
|
||||
pub fn estimate_grid_size(complexity: f64, base_size: usize, max_size: usize) -> usize {
|
||||
let adaptive_factor = (complexity * 10.0).max(1.0).min(5.0);
|
||||
let estimated_size = (base_size as f64 * adaptive_factor) as usize;
|
||||
estimated_size.min(max_size).max(base_size)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, feature = "disabled_tests"))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_kan_config_creation() {
|
||||
let config = KANConfig::new(vec![3, 10, 5, 1]);
|
||||
assert_eq!(config.architecture, vec![3, 10, 5, 1]);
|
||||
assert_eq!(config.grid_size, 20);
|
||||
assert_eq!(config.basis_function_type, BasisFunctionType::BSpline);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kan_config_validation() {
|
||||
// Valid config
|
||||
let valid_config = KANConfig::new(vec![2, 5, 1]);
|
||||
assert!(valid_config.validate().is_ok());
|
||||
|
||||
// Invalid architecture (too short)
|
||||
let invalid_config = KANConfig::new(vec![5]);
|
||||
assert!(invalid_config.validate().is_err());
|
||||
|
||||
// Invalid grid size
|
||||
let mut invalid_grid = KANConfig::new(vec![2, 5, 1]);
|
||||
invalid_grid.grid_size = 2;
|
||||
assert!(invalid_grid.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_builder_methods() {
|
||||
let config = KANConfig::new(vec![2, 8, 1])
|
||||
.with_grid_size(30)
|
||||
.with_spline_order(4)
|
||||
.with_basis_function(BasisFunctionType::Chebyshev)
|
||||
.with_basis_range(-1.0, 1.0)
|
||||
.with_adaptive_grid(true, 0.02, 80)
|
||||
.with_pruning(true, 0.6)
|
||||
.with_symbolic_regression(true, 8);
|
||||
|
||||
assert_eq!(config.grid_size, 30);
|
||||
assert_eq!(config.spline_order, 4);
|
||||
assert_eq!(config.basis_function_type, BasisFunctionType::Chebyshev);
|
||||
assert_eq!(config.basis_range, (-1.0, 1.0));
|
||||
assert!(config.adaptive_grid);
|
||||
assert_eq!(config.adaptation_threshold, 0.02);
|
||||
assert_eq!(config.max_grid_size, 80);
|
||||
assert!(config.enable_pruning);
|
||||
assert_eq!(config.sparsity_target, 0.6);
|
||||
assert!(config.enable_symbolic);
|
||||
assert_eq!(config.symbolic_max_depth, 8);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user