tensorlogic-trustformers

Transformer architectures as TensorLogic einsum graphs

This crate provides implementations of transformer components (self-attention, multi-head attention, feed-forward networks) as einsum operations that compile to TensorLogic IR and execute on any TensorLogic backend.

Features

• ✅ Self-Attention - Scaled dot-product attention as einsum operations
• ✅ Multi-Head Attention - Parallel attention heads with automatic head splitting/merging
• ✅ Feed-Forward Networks - Position-wise FFN with configurable activations (GELU, ReLU, etc.)
• ✅ Gated FFN - GLU-style gated feed-forward networks
• ✅ Position Encodings - Sinusoidal, learned, and relative position encodings
• ✅ Layer Normalization - Standard LayerNorm and RMSNorm implementations
• ✅ Encoder Layers - Complete transformer encoder layers with pre/post-norm variants
• ✅ Decoder Layers - Complete transformer decoder layers with masked self-attention
• ✅ Encoder/Decoder Stacks - Multi-layer transformer stacks with flexible configuration
• ✅ Rule-Based Attention - Logical rules guiding attention patterns (hard/soft/gated)
• ✅ Sparse Attention - Efficient attention for long sequences (strided, local, block-sparse)
• ✅ Utility Functions - Parameter counting, FLOP calculations, model presets
• ✅ Gradient Checkpointing - Memory-efficient training with uniform/selective/dynamic strategies
• ✅ KV-Cache - Efficient autoregressive inference with 10-1000x speedup
• ✅ Performance Benchmarks - Criterion-based benchmark suite with HTML reports
• ✅ Type-Safe Configuration - Builder pattern with validation
• ✅ Einsum-Native - All operations expressed as einsum for maximum flexibility
• ✅ Zero Warnings - Strict code quality enforcement
• ✅ 229 Tests - Comprehensive test coverage (100% passing)

Quick Start

use tensorlogic_trustformers::{
    AttentionConfig, SelfAttention, MultiHeadAttention,
    FeedForwardConfig, FeedForward,
};
use tensorlogic_ir::EinsumGraph;

// Configure and build self-attention
let attn_config = AttentionConfig::new(512, 8).unwrap();
let self_attn = SelfAttention::new(attn_config).unwrap();

let mut graph = EinsumGraph::new();
graph.add_tensor("Q");
graph.add_tensor("K");
graph.add_tensor("V");

let outputs = self_attn.build_attention_graph(&mut graph).unwrap();

// Configure multi-head attention
let mha_config = AttentionConfig::new(512, 8).unwrap();
let mha = MultiHeadAttention::new(mha_config).unwrap();

// Configure feed-forward network
let ffn_config = FeedForwardConfig::new(512, 2048)
    .with_activation("gelu")
    .with_dropout(0.1);
let ffn = FeedForward::new(ffn_config).unwrap();

Architecture

Self-Attention Formula

Attention(Q, K, V) = softmax(QK^T / √d_k) V

Einsum breakdown:

1. Query-Key scores: einsum("bqd,bkd->bqk", Q, K)
2. Scale: scores / sqrt(d_k)
3. Softmax: softmax(scores, axis=-1)
4. Attention-Value: einsum("bqk,bkv->bqv", attn, V)

Where:

• b = batch dimension
• q = query sequence length
• k = key sequence length
• d = model dimension
• v = value dimension

Multi-Head Attention

Multi-head attention splits the model dimension into parallel attention heads:

1. Reshape: [B, S, D] -> [B, H, S, D_k] where D_k = D/H
2. Attention per head: einsum("bhqd,bhkd->bhqk", Q, K)
3. Scale and softmax
4. Apply to values: einsum("bhqk,bhkv->bhqv", attn, V)
5. Concatenate heads: [B, H, S, D_k] -> [B, S, D]

Feed-Forward Network

Position-wise feed-forward network with two linear transformations:

FFN(x) = activation(xW1 + b1)W2 + b2

Einsum notation:

1. First linear: einsum("bsd,df->bsf", x, W1)
2. Activation: activation(h1) (GELU, ReLU, etc.)
3. Second linear: einsum("bsf,fd->bsd", h2, W2)

Where:

• d = d_model
• f = d_ff (typically 4 * d_model)

Configuration

Attention Configuration

use tensorlogic_trustformers::AttentionConfig;

let config = AttentionConfig::new(512, 8)?
    .with_causal(true)      // Enable causal masking
    .with_dropout(0.1);      // Set dropout probability

assert_eq!(config.d_model, 512);
assert_eq!(config.n_heads, 8);
assert_eq!(config.d_k, 64);  // Automatically computed

Feed-Forward Configuration

use tensorlogic_trustformers::FeedForwardConfig;

let config = FeedForwardConfig::new(512, 2048)
    .with_activation("gelu")  // or "relu", "silu", etc.
    .with_dropout(0.1);

assert_eq!(config.d_model, 512);
assert_eq!(config.d_ff, 2048);

Complete Transformer Layer

use tensorlogic_trustformers::TransformerLayerConfig;

let config = TransformerLayerConfig::new(512, 8, 2048)?
    .with_pre_norm(true);   // Use pre-layer normalization

assert!(config.validate().is_ok());

Graph Building

Self-Attention Graph

use tensorlogic_trustformers::SelfAttention;
use tensorlogic_ir::EinsumGraph;

let attn = SelfAttention::new(config)?;
let mut graph = EinsumGraph::new();

// Add input tensors (Q, K, V)
graph.add_tensor("Q");  // [batch, seq, d_model]
graph.add_tensor("K");  // [batch, seq, d_model]
graph.add_tensor("V");  // [batch, seq, d_model]

// Build attention graph
let outputs = attn.build_attention_graph(&mut graph)?;
// outputs[0] = attention output [batch, seq, d_model]

Multi-Head Attention Graph

use tensorlogic_trustformers::MultiHeadAttention;

let mha = MultiHeadAttention::new(config)?;
let mut graph = EinsumGraph::new();

graph.add_tensor("Q");
graph.add_tensor("K");
graph.add_tensor("V");

let outputs = mha.build_mha_graph(&mut graph)?;

Feed-Forward Network Graph

use tensorlogic_trustformers::FeedForward;

let ffn = FeedForward::new(config)?;
let mut graph = EinsumGraph::new();

// Add input tensors
graph.add_tensor("x");   // [batch, seq, d_model]
graph.add_tensor("W1");  // [d_model, d_ff]
graph.add_tensor("b1");  // [d_ff]
graph.add_tensor("W2");  // [d_ff, d_model]
graph.add_tensor("b2");  // [d_model]

let outputs = ffn.build_ffn_graph(&mut graph)?;

Advanced Features

Gated Feed-Forward Network (GLU)

GLU-style networks use element-wise gating for improved capacity:

use tensorlogic_trustformers::GatedFeedForward;

let glu = GatedFeedForward::new(config)?;
let mut graph = EinsumGraph::new();

graph.add_tensor("x");
graph.add_tensor("W_gate");
graph.add_tensor("W_value");
graph.add_tensor("W_out");

let outputs = glu.build_glu_graph(&mut graph)?;

Formula: GLU(x) = σ(xW_gate) ⊙ activation(xW_value) W_out

Integration with TensorLogic

The einsum graphs produced by this crate integrate seamlessly with the TensorLogic ecosystem:

Compilation

use tensorlogic_compiler::CompilerContext;

let mut ctx = CompilerContext::new();
// Compile TLExpr rules that use transformer operations

Execution

use tensorlogic_scirs_backend::Scirs2Executor;

let executor = Scirs2Executor::new();
// Execute the transformer graph on SciRS2 backend

Optimization

use tensorlogic_ir::graph::optimization::optimize_graph;

let stats = optimize_graph(&mut graph)?;
// Apply dead code elimination, CSE, etc.

Design Philosophy

This crate follows core TensorLogic principles:

1. Backend Independence: Same graph works on CPU, GPU, TPU
2. Einsum-Native: Clear mathematical semantics
3. Composability: Mix transformer layers with logical rules
4. Type Safety: Compile-time dimension checking where possible
5. Zero Cost Abstractions: No runtime overhead

Examples

See the examples directory for complete examples:

• 01_basic_encoder.rs - Basic transformer encoder usage
• 02_trustformers_integration.rs - TrustformeRS integration
• 03_rule_based_attention.rs - Rule-based attention patterns
• 04_sparse_attention.rs - Sparse attention for long sequences
• 05_gradient_checkpointing.rs - Memory-efficient training strategies
• 06_kv_cache_inference.rs - Fast autoregressive generation with KV-cache

Testing

Run the test suite:

cargo nextest run -p tensorlogic-trustformers

All 229 tests should pass with zero warnings.

Benchmarking

Run performance benchmarks:

cargo bench --bench model_benchmarks

This will generate HTML reports in target/criterion/ with detailed performance metrics.

Performance

The einsum-based approach enables:

• Operation Fusion: Compiler can fuse consecutive operations
• Memory Efficiency: Minimal intermediate tensors
• Parallelization: Natural SIMD/GPU mapping
• Optimization: Graph-level optimizations

Roadmap

See TODO.md for the development roadmap. Current status: 100% complete 🎉

Completed ✅

• Self-attention as einsum
• Multi-head attention
• Feed-forward networks (standard + gated GLU)
• Position encodings (sinusoidal, learned, relative, RoPE, ALiBi)
• Layer normalization (LayerNorm + RMSNorm)
• Transformer encoder layers (pre-norm + post-norm)
• Transformer decoder layers (pre-norm + post-norm)
• Encoder/decoder stacks with position encoding
• Rule-based attention patterns (hard/soft/gated)
• Sparse attention patterns (strided, local, block-sparse, global-local)
• Gradient checkpointing (uniform, selective, dynamic)
• KV-cache for efficient inference (10-1000x speedup)
• TrustformeRS integration (bidirectional conversion)
• Utility functions (parameter counting, FLOP calculations, presets)
• Performance benchmarking suite (Criterion)
• Configuration system with validation
• Error handling with IrError conversion
• 229 comprehensive tests (100% passing, zero warnings)
• 6 complete examples

Future Enhancements 📋

• Vision transformers (ViT)
• Flash Attention integration
• Pre-trained model weight import
• Advanced pattern composition
• GPU-specific optimizations
• Speculative decoding
• Quantization support

References

• Attention Is All You Need - Original transformer paper
• Tensor Logic Paper - TensorLogic framework
• Einsum Documentation - Einsum notation

License

This crate is part of the TensorLogic project and is licensed under Apache-2.0.

New Features in v0.1.0

Position Encodings

Three types of position encodings for sequence modeling:

use tensorlogic_trustformers::{PositionEncodingConfig, SinusoidalPositionEncoding};

// Sinusoidal (fixed) encoding
let config = PositionEncodingConfig::sinusoidal(512, 2048);
let pe = SinusoidalPositionEncoding::new(config).unwrap();

// Learned position embeddings
let config = PositionEncodingConfig::learned(512, 2048);
let pe = LearnedPositionEncoding::new(config).unwrap();

// Relative position encoding
let config = PositionEncodingConfig::relative(512, 32, 128);
let pe = RelativePositionEncoding::new(config).unwrap();

Layer Normalization

Standard LayerNorm and efficient RMSNorm:

use tensorlogic_trustformers::{LayerNormConfig, LayerNorm, RMSNorm};

// Standard layer normalization
let config = LayerNormConfig::new(512).with_eps(1e-6);
let ln = LayerNorm::new(config).unwrap();

// RMS normalization (more efficient)
let rms = RMSNorm::new(config).unwrap();

Complete Transformer Layers

Full encoder and decoder layers with residual connections:

use tensorlogic_trustformers::{EncoderLayerConfig, EncoderLayer};

// Encoder layer with pre-normalization
let config = EncoderLayerConfig::new(512, 8, 2048)?
    .with_pre_norm(true)
    .with_dropout(0.1);
let encoder = EncoderLayer::new(config)?;

// Decoder layer with causal masking
let decoder_config = DecoderLayerConfig::new(512, 8, 2048)?;
let decoder = DecoderLayer::new(decoder_config)?;

Transformer Stacks

Multi-layer transformer architectures:

use tensorlogic_trustformers::{EncoderStackConfig, EncoderStack};

// 6-layer transformer encoder
let config = EncoderStackConfig::new(6, 512, 8, 2048, 1024)?
    .with_dropout(0.1)
    .with_final_layer_norm(true);
let encoder_stack = EncoderStack::new(config)?;

// Build complete encoder graph
let mut graph = EinsumGraph::new();
graph.add_tensor("input");
let outputs = encoder_stack.build_encoder_stack_graph(&mut graph)?;

Rule-Based Attention

Integrate logical rules with attention mechanisms:

use tensorlogic_trustformers::{RuleAttentionConfig, RuleBasedAttention};
use tensorlogic_trustformers::rule_attention::patterns;

// Hard constraint: only attend where rule is satisfied
let base_attn = AttentionConfig::new(512, 8)?;
let config = RuleAttentionConfig::hard(base_attn);
let rule = patterns::syntactic_dependency("head", "dep");
let attn = RuleBasedAttention::new(config)?.with_rule(rule);

// Soft constraint: bias attention towards rule-satisfying positions
let config = RuleAttentionConfig::soft(base_attn, 0.7);

// Gated: interpolate between content and rule attention
let config = RuleAttentionConfig::gated(base_attn, 0.5);

Gradient Checkpointing

Memory-efficient training for large models:

use tensorlogic_trustformers::{CheckpointConfig, EncoderStackConfig};

// Create a large model
let config = EncoderStackConfig::new(12, 768, 12, 3072, 512)?;

// Uniform checkpointing: checkpoint every 2 layers
let checkpoint = CheckpointConfig::uniform(2);
println!("Memory savings: {:.1}%", checkpoint.memory_savings(12) * 100.0);
println!("Compute overhead: {:.2}x", checkpoint.compute_overhead(12));

// Selective checkpointing: checkpoint specific layers
let checkpoint = CheckpointConfig::selective(vec![0, 3, 6, 9]);

// Dynamic checkpointing: automatically balance memory vs. compute
let checkpoint = CheckpointConfig::dynamic(12, 0.3)?; // Target 30% memory usage

// Customize what to checkpoint
let checkpoint = CheckpointConfig::uniform(2)
    .with_checkpoint_attention(true)   // Checkpoint attention
    .with_checkpoint_ffn(false);       // Don't checkpoint FFN

Benefits:

• 50-80% memory savings depending on strategy
• 1.1-1.3x compute overhead (modest increase)
• Train larger models or use bigger batch sizes
• Three strategies: uniform, selective, dynamic

KV-Cache for Fast Inference

Enable efficient autoregressive generation with dramatic speedups:

use tensorlogic_trustformers::{KVCache, KVCacheConfig};

// Create cache for 12-layer model (GPT-2 small)
let mut cache = KVCache::new(12, 12, 64);

// During autoregressive generation
for step in 0..100 {
    // Compute keys/values only for new token
    let keys = compute_keys_for_new_token();   // [batch, heads, 1, dim]
    let values = compute_values_for_new_token(); // [batch, heads, 1, dim]

    // Update cache for all layers
    for layer_idx in 0..12 {
        cache.update_layer(layer_idx, keys.clone(), values.clone())?;
    }

    // Retrieve cached keys/values for attention
    let (all_keys, all_values) = cache.get_layer(0)?;

    // Compute attention only over new position
    // ... (attention computation using cached K,V)

    cache.next_step();
}

// Monitor cache usage
let stats = cache.stats();
println!("{}", stats.summary());
// CacheStats:
//   Layers: 12
//   Seq len: 100
//   Memory: 7.0/4608.0 MB (0.2%)
//   Step: 100
//   Enabled: true

Performance Impact:

• 10-1000x speedup depending on sequence length
• Linear speedup with sequence length: 100 tokens = 100x faster
• Minimal memory cost: ~2-10 MB for typical models
• Essential for production text generation

Configuration Options:

// Custom cache configuration
let config = KVCacheConfig::new(24, 16, 64)  // GPT-2 large
    .with_max_seq_len(4096)    // Support longer contexts
    .with_max_batch_size(64)   // Larger batch inference
    .with_enabled(true);       // Enable/disable dynamically

let cache = KVCache::from_config(config)?;

// Memory estimation
println!("Max memory: {:.1} MB", config.memory_usage_mb());

---

Sparse Attention

Efficient attention for long sequences:

use tensorlogic_trustformers::{SparseAttentionConfig, SparseAttention, LocalAttention};

// Strided sparse attention (attend every k-th position)
let base_attn = AttentionConfig::new(512, 8)?;
let config = SparseAttentionConfig::strided(base_attn, 4)?;
let sparse = SparseAttention::new(config)?;

// Local windowed attention
let config = SparseAttentionConfig::local(base_attn, 128)?;
let sparse = SparseAttention::new(config)?;

// Or use dedicated LocalAttention for efficiency
let local = LocalAttention::new(base_attn, 64)?;
println!("Memory savings: {:.1}%", local.memory_savings(1024) * 100.0);

Utility Functions

Helper functions for model analysis:

use tensorlogic_trustformers::utils::{encoder_stack_stats, presets};

// Get model statistics
let config = presets::gpt2_small();
let stats = encoder_stack_stats(&config);
println!("{}", stats.summary());
// Output: ModelStats:
//   Total params: 117.00M
//   Trainable: 117.00M
//   Layers: 12
//   d_model: 768
//   Memory: 468 MB

// Use preset configurations
let gpt2 = presets::gpt2_small();
let bert = presets::bert_base();
let (encoder, decoder) = presets::transformer_base();

---

Status: 🎉 Production Ready (v0.1.0-alpha.2) **Last Updated: 2025-12-16 Tests: 229/229 passing (100%) Examples: 6 comprehensive examples Benchmarks: Criterion suite with HTML reports Features: Complete transformer implementation with optimizations Part of: TensorLogic Ecosystem