Fock-Mode Attention (FMA) - A Memory-Efficient Transformer Architecture

🎯 TL;DR

📚 Table of Contents

🧠 Theory: What is Fock-Mode Attention? {#theory}

The Problem with Standard Attention

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

Fock-Mode Inspiration (Quantum Mechanics)

Mathematical Formulation

Output = softmax(Q K^T) V
Size: (×ばつN) @ (×ばつD) = O(N2)

1. Emission: G = softmax(X W_g) (×ばつM)
2. Projection: S = G^T Z (×ばつD)
3. Mode mixing: S' = MLP(S) (×ばつD)
4. Absorption: H = softmax(X W_h) (×ばつM)
5. Output: Y = H S' (×ばつD)
Total complexity: O(×ばつM) + O(×ばつD2)

Analogy: Hub-and-Spoke Communication

City1 ←→ City2
 ↕ ↕
City3 ←→ City4
... (N2 connections)

Cities → [Hub1, Hub2, ..., Hub16] → Cities
(N→M) (M mixing) (M→N)

🏗️ Architecture Overview {#architecture}

Core Components

Input Sequence (×ばつD)
 ↓
┌──────────────────────┐
│ FockModeAttention │
│ ┌────────────────┐ │
│ │ 1. Emission │ │ G = softmax(X W_g) (N→M)
│ │ 2. Token→Mode │ │ S = G^T Z
│ │ 3. Mode Mix │ │ S' = MLP(S)
│ │ 4. Mode→Token │ │ Y = H S'
│ │ 5. Absorption │ │ H = softmax(X W_h)
│ └────────────────┘ │
└──────────────────────┘
 ↓
 LayerNorm + FFN
 ↓
 Output (×ばつD)

Full Model: FMAEncoderModel

FMAEncoderModel
├── Token Embedding (vocab_size → d_model)
├── Positional Embedding (max_len → d_model)
├── N ×ばつ FMAEncoderBlock
│ ├── FockModeAttention (d_model, num_modes)
│ ├── LayerNorm
│ ├── FeedForward (4×ばつd_model)
│ └── LayerNorm
└── Classification Head (d_model → num_classes)

💻 Implementation Details {#implementation}

File Structure

FSNN/
├── core/
│ ├── attention.py # FockModeAttention + FastFockModeAttention
│ └── layers.py # FMAEncoderBlock
├── models/
│ ├── fma_model.py # Full FMA model
│ └── baseline_model.py # Standard Transformer (for comparison)
├── training/
│ ├── train_tiny.py # Quick demo (synthetic data)
│ └── train_distill.py # Real KD (IMDb + BERT-Tiny)
├── experiments/
│ ├── benchmark_attention_fast.py # Speed benchmark
│ ├── benchmark_memory.py # Memory benchmark
│ └── full_comparison.py # Model comparison
├── data/
│ └── synthetic.py # Data generation
├── test_model.py # Inference on text prompts
└── checkpoints/ # Saved models

Key Optimizations

# Uses torch.matmul (multiple kernel launches)
S = torch.matmul(g.transpose(1, 2), z)
Y = torch.matmul(h, S_mixed)

# Uses einsum (fused kernels)
S = torch.einsum("bnm,bnd->bmd", g, z)
Y = torch.einsum("bnm,bmd->bnd", h, S_mixed)
# Conv1d for mode mixing (faster than Linear)
self.mode_conv1 = nn.Conv1d(d_inner, 4*d_inner, kernel_size=1)

📊 Performance Analysis {#performance}

Hardware: RTX 4050 6GB Laptop GPU

1. Speed Benchmark (B=32, N=512, D=256)

2. Memory Benchmark (B=1, D=256, M=16)

Memory reduction ≈ 1 - (×ばつH)/N
As N increases → reduction approaches 100%!

3. Maximum Sequence Length (6GB GPU)

4. Model Size

⚖️ Pros & Cons {#pros-cons}

✅ Advantages

❌ Disadvantages

When to Use FMA vs Standard Attention

🚀 Usage & Setup {#usage}

Environment Setup

# Create environment
conda create -n fsnn python=3.10 -y
conda activate fsnn
# Install dependencies
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install transformers datasets accelerate

Quick Start: Train & Test

# 1. Train on synthetic data (30 seconds)
python -m training.train_tiny
# 2. Train with real data (10-15 minutes)
python -m training.train_distill
# 3. Test with text prompts
python test_model.py

Example: Text Classification

from models.fma_model import FMAEncoderModel
from transformers import AutoTokenizer
import torch
# Load model
tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-tiny")
model = FMAEncoderModel(
 vocab_size=tokenizer.vocab_size,
 d_model=128,
 num_layers=2,
 num_modes=32,
 max_len=128,
 num_classes=2
).cuda()
# Load checkpoint
model.load_state_dict(torch.load("checkpoints/student_fma_distilled.pt"))
model.eval()
# Inference
text = "This movie was amazing! I loved it."
inputs = tokenizer(text, return_tensors="pt", max_length=128, truncation=True)
with torch.no_grad():
 logits = model(inputs['input_ids'].cuda())
 prediction = torch.argmax(logits, dim=-1)
 
print(f"Sentiment: {'Positive' if prediction.item() == 1 else 'Negative'}")

📈 Benchmarks {#benchmarks}

Run All Benchmarks

# Speed comparison
python -m experiments.benchmark_attention_fast
# Memory comparison 
python -m experiments.benchmark_memory
# Full model comparison
python -m experiments.full_comparison

Sample Output

=== Memory Benchmark ===
Seq Len | Standard (MB) | FMA (MB) | Savings (MB) | Reduction %
--------|---------------|----------|--------------|-------------
 128 | 10.02 | 8.67 | 1.35 | 13.5%
 2048 | 36.67 | 29.72 | 6.95 | 18.9%
 4096 | ~150.0 | ~60.0 | ~90.0 | ~60.0%
KEY INSIGHT: Memory savings GROW with sequence length

🔬 Knowledge Distillation Pipeline

Setup

Results

Teacher Accuracy: 49% (frozen, not trained on IMDb)
Student Accuracy: 70-85% (after 10 epochs KD)

🎓 Theory: Why Does This Work?

Three Key Insights

🔮 Future Work {#future-work}

Immediate Improvements

Long-term Research

📝 Citation

@software{fock_mode_attention_2024,
 title = {Fock-Mode Attention: A Memory-Efficient Transformer Architecture},
 author = {[Your Name]},
 year = {2024},
 url = {https://github.com/yourusername/FSNN},
 note = {Implementation with Knowledge Distillation on RTX 4050 6GB}
}

🤝 Acknowledgments

📄 License

🎯 Key Takeaways