200 lines
6.6 KiB
Rust
200 lines
6.6 KiB
Rust
//! Mamba (State Space Model) Demonstration
|
|
//!
|
|
//! This example demonstrates the Mamba architecture implementation,
|
|
//! showcasing its linear complexity and integration with transformer components.
|
|
|
|
use rtx_tensor::{Device, Tensor};
|
|
use rtx_transformers::layers::{MambaBlock, MambaConfig, MambaTransformer, MambaTransformerBlock};
|
|
use std::time::Instant;
|
|
|
|
/// Demonstrates basic Mamba block functionality
|
|
async fn demo_mamba_block() -> Result<(), Box<dyn std::error::Error>> {
|
|
println!("🔬 Mamba Block Demo");
|
|
println!("===================");
|
|
|
|
let device = Device::cpu();
|
|
|
|
// Create Mamba configuration
|
|
let config = MambaConfig::new(512, 16, 4); // d_model=512, d_state=16, d_conv=4
|
|
println!(
|
|
"Configuration: d_model={}, d_state={}, d_conv={}",
|
|
config.d_model, config.d_state, config.d_conv
|
|
);
|
|
|
|
// Initialize Mamba block
|
|
let mamba = MambaBlock::new(config.clone(), &device)?;
|
|
println!(
|
|
"✅ Mamba block created with {} parameters",
|
|
mamba.parameters().len()
|
|
);
|
|
|
|
// Test with different sequence lengths
|
|
let sequence_lengths = vec![10, 50, 100, 200];
|
|
println!("\n📊 Linear Complexity Demonstration:");
|
|
|
|
for seq_len in sequence_lengths {
|
|
let batch_size = 2;
|
|
let input = Tensor::randn([batch_size, seq_len, config.d_model], &device)?;
|
|
|
|
let start = Instant::now();
|
|
let output = mamba.forward(&input)?;
|
|
let elapsed = start.elapsed();
|
|
|
|
println!(
|
|
"Seq length {}: {:.2}ms (output shape: {:?})",
|
|
seq_len,
|
|
elapsed.as_millis(),
|
|
output.output.shape().dims()
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Demonstrates hybrid Mamba-Transformer architecture
|
|
async fn demo_mamba_transformer() -> Result<(), Box<dyn std::error::Error>> {
|
|
println!("\n🏗️ Mamba Transformer Demo");
|
|
println!("==========================");
|
|
|
|
let device = Device::cpu();
|
|
|
|
// Create hybrid transformer block
|
|
let d_model = 256;
|
|
let ff_dim = 1024;
|
|
let mamba_config = MambaConfig::new(d_model, 16, 4);
|
|
|
|
let hybrid_block = MambaTransformerBlock::new(d_model, ff_dim, mamba_config, &device)?;
|
|
println!("✅ Hybrid Mamba-Transformer block created");
|
|
|
|
// Test forward pass
|
|
let batch_size = 1;
|
|
let seq_len = 20;
|
|
let input = Tensor::randn([batch_size, seq_len, d_model], &device)?;
|
|
|
|
let start = Instant::now();
|
|
let output = hybrid_block.forward(&input)?;
|
|
let elapsed = start.elapsed();
|
|
|
|
println!("Forward pass: {:.2}ms", elapsed.as_millis());
|
|
println!("Input shape: {:?}", input.shape().dims());
|
|
println!("Output shape: {:?}", output.shape().dims());
|
|
|
|
// Verify residual connections work
|
|
let output_data = output.to_vec()?;
|
|
let input_data = input.to_vec()?;
|
|
let are_different = output_data
|
|
.iter()
|
|
.zip(input_data.iter())
|
|
.any(|(&a, &b)| (a - b).abs() > 1e-6);
|
|
|
|
println!("Residual connection working: {}", are_different);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Demonstrates full Mamba transformer model
|
|
async fn demo_full_transformer() -> Result<(), Box<dyn std::error::Error>> {
|
|
println!("\n🧠 Full Mamba Transformer Demo");
|
|
println!("===============================");
|
|
|
|
let device = Device::cpu();
|
|
|
|
// Model configuration
|
|
let vocab_size = 1000;
|
|
let d_model = 128;
|
|
let num_layers = 3;
|
|
let max_seq_len = 100;
|
|
|
|
// Create full transformer
|
|
let transformer = MambaTransformer::new(vocab_size, d_model, num_layers, max_seq_len, &device)?;
|
|
|
|
println!("✅ Full Mamba Transformer created:");
|
|
println!(" - Vocabulary size: {}", vocab_size);
|
|
println!(" - Model dimension: {}", d_model);
|
|
println!(" - Number of layers: {}", num_layers);
|
|
println!(" - Max sequence length: {}", max_seq_len);
|
|
|
|
// Test inference
|
|
let batch_size = 1;
|
|
let seq_len = 15;
|
|
let input_ids = Tensor::randint(0, vocab_size as i32, &[batch_size, seq_len], &device)?;
|
|
|
|
let start = Instant::now();
|
|
let logits = transformer.forward(&input_ids)?;
|
|
let elapsed = start.elapsed();
|
|
|
|
println!("Inference time: {:.2}ms", elapsed.as_millis());
|
|
println!("Input shape: {:?}", input_ids.shape().dims());
|
|
println!("Output logits shape: {:?}", logits.shape().dims());
|
|
|
|
// Verify output is reasonable
|
|
let logits_data = logits.to_vec()?;
|
|
let max_logit = logits_data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
|
let min_logit = logits_data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
|
|
|
|
println!("Logit range: [{:.3}, {:.3}]", min_logit, max_logit);
|
|
println!(
|
|
"All logits finite: {}",
|
|
logits_data.iter().all(|&x| x.is_finite())
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Compares Mamba efficiency with theoretical attention complexity
|
|
async fn demo_efficiency_comparison() -> Result<(), Box<dyn std::error::Error>> {
|
|
println!("\n⚡ Efficiency Comparison");
|
|
println!("========================");
|
|
|
|
let device = Device::cpu();
|
|
let d_model = 256;
|
|
let config = MambaConfig::new(d_model, 16, 4);
|
|
let mamba = MambaBlock::new(config, &device)?;
|
|
|
|
println!("Comparing Mamba (linear) vs theoretical Attention (quadratic):");
|
|
println!("Seq Length | Mamba Time (ms) | Theoretical Attention FLOPs");
|
|
println!("-----------|-----------------|-----------------------------");
|
|
|
|
for seq_len in [50, 100, 200, 400, 800] {
|
|
let input = Tensor::randn([1, seq_len, d_model], &device)?;
|
|
|
|
let start = Instant::now();
|
|
let _output = mamba.forward(&input)?;
|
|
let mamba_time = start.elapsed().as_millis();
|
|
|
|
// Theoretical attention FLOPs (simplified): O(seq_len^2 * d_model)
|
|
let attention_flops = (seq_len * seq_len * d_model) as f64 / 1e6; // In millions
|
|
|
|
println!(
|
|
"{:10} | {:15} | {:.2}M FLOPs",
|
|
seq_len, mamba_time, attention_flops
|
|
);
|
|
}
|
|
|
|
println!("\n📈 Mamba scales linearly O(n) while attention scales quadratically O(n²)");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
println!("🐍 Mamba: Linear-Time Sequence Modeling Demo");
|
|
println!("===========================================");
|
|
println!("State Space Models for Transformers");
|
|
println!("Paper: https://arxiv.org/abs/2312.00752\n");
|
|
|
|
demo_mamba_block().await?;
|
|
demo_mamba_transformer().await?;
|
|
demo_full_transformer().await?;
|
|
demo_efficiency_comparison().await?;
|
|
|
|
println!("\n🎉 All Mamba demonstrations completed successfully!");
|
|
println!("Key advantages demonstrated:");
|
|
println!("- ✅ Linear complexity O(n) vs attention's O(n²)");
|
|
println!("- ✅ Selective state space modeling");
|
|
println!("- ✅ Integration with transformer architecture");
|
|
println!("- ✅ Efficient caching for inference");
|
|
|
|
Ok(())
|
|
}
|