Files

197 lines
6.6 KiB
Markdown

# REFRAG Pipeline Example
> **Compress-Sense-Expand Architecture for ~30x RAG Latency Reduction**
This example demonstrates the REFRAG (Rethinking RAG) framework from [arXiv:2509.01092](https://arxiv.org/abs/2509.01092) using ruvector as the underlying vector store.
## Overview
Traditional RAG systems return text chunks that must be tokenized and processed by the LLM. REFRAG instead stores pre-computed "representation tensors" and uses a lightweight policy network to decide whether to return:
- **COMPRESS**: The tensor representation (directly injectable into LLM context)
- **EXPAND**: The original text (for cases where full context is needed)
## Architecture
```
┌─────────────────────────────────────────────────────────────────┐
│ REFRAG Pipeline │
├─────────────────────────────────────────────────────────────────┤
│ │
│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │
│ │ COMPRESS │ │ SENSE │ │ EXPAND │ │
│ │ Layer │───▶│ Layer │───▶│ Layer │ │
│ └──────────────┘ └──────────────┘ └──────────────┘ │
│ │
│ Binary tensor Policy network Dimension projection │
│ storage with decides COMPRESS (768 → 4096 dims) │
│ zero-copy access vs EXPAND │
│ │
└─────────────────────────────────────────────────────────────────┘
```
### Compress Layer (`compress.rs`)
Stores representation tensors in binary format with multiple compression strategies:
| Strategy | Compression | Use Case |
|----------|-------------|----------|
| `None` | 1x | Maximum precision |
| `Float16` | 2x | Good balance |
| `Int8` | 4x | Memory constrained |
| `Binary` | 32x | Extreme compression |
### Sense Layer (`sense.rs`)
Policy network that decides the response type for each retrieved chunk:
| Policy | Latency | Description |
|--------|---------|-------------|
| `ThresholdPolicy` | ~2μs | Cosine similarity threshold |
| `LinearPolicy` | ~5μs | Single layer classifier |
| `MLPPolicy` | ~15μs | Two-layer neural network |
### Expand Layer (`expand.rs`)
Projects tensors to target LLM dimensions when needed:
| Source | Target | LLM |
|--------|--------|-----|
| 768 | 4096 | LLaMA-3 8B |
| 768 | 8192 | LLaMA-3 70B |
| 1536 | 8192 | GPT-4 |
## Quick Start
```bash
# Run the demo
cargo run --bin refrag-demo
# Run benchmarks (use release for accurate measurements)
cargo run --bin refrag-benchmark --release
```
## Usage
### Basic Usage
```rust
use refrag_pipeline_example::{RefragStore, RefragEntry};
// Create REFRAG-enabled store
let store = RefragStore::new(384, 768)?;
// Insert with representation tensor
let entry = RefragEntry::new("doc_1", search_vector, "The quick brown fox...")
.with_tensor(tensor_bytes, "llama3-8b");
store.insert(entry)?;
// Standard search (text only)
let results = store.search(&query, 10)?;
// Hybrid search (policy-based COMPRESS/EXPAND)
let results = store.search_hybrid(&query, 10, Some(0.85))?;
for result in results {
match result.response_type {
RefragResponseType::Compress => {
println!("Tensor: {} dims", result.tensor_dims.unwrap());
}
RefragResponseType::Expand => {
println!("Text: {}", result.content.unwrap());
}
}
}
```
### Custom Configuration
```rust
use refrag_pipeline_example::{
RefragStoreBuilder,
PolicyNetwork,
ExpandLayer,
};
let store = RefragStoreBuilder::new()
.search_dimensions(384)
.tensor_dimensions(768)
.target_dimensions(4096)
.compress_threshold(0.85) // Higher = more COMPRESS
.auto_project(true)
.policy(PolicyNetwork::mlp(768, 32, 0.85))
.expand_layer(ExpandLayer::for_roberta())
.build()?;
```
### Response Format
REFRAG search returns a hybrid response format:
```json
{
"results": [
{
"id": "doc_1",
"score": 0.95,
"response_type": "EXPAND",
"content": "The quick brown fox...",
"policy_confidence": 0.92
},
{
"id": "doc_2",
"score": 0.88,
"response_type": "COMPRESS",
"tensor_b64": "base64_encoded_float32_array...",
"tensor_dims": 4096,
"alignment_model_id": "llama3-8b",
"policy_confidence": 0.97
}
]
}
```
## Performance
### Latency Breakdown
| Component | Latency |
|-----------|---------|
| Vector search (HNSW) | 100-500μs |
| Policy decision | 1-50μs |
| Tensor decompression | 1-10μs |
| Projection (optional) | 10-100μs |
| **Total** | **~150-700μs** |
### Comparison to Traditional RAG
| Operation | Traditional | REFRAG |
|-----------|-------------|--------|
| Text tokenization | 1-5ms | N/A |
| LLM context prep | 5-20ms | ~100μs |
| Network transfer | 10-50ms | ~1-5ms |
| **Speedup** | - | **10-30x** |
## Why REFRAG Works for RuVector
1. **Rust/WASM**: Python implementations suffer from loop overhead. RuVector runs the policy in SIMD-optimized Rust (<50μs decisions).
2. **Edge Deployment**: The WASM build can serve as a "Smart Context Compressor" in the browser, sending only necessary tokens/tensors to the server LLM.
3. **Zero-Copy**: Using `rkyv` serialization enables direct memory access to tensors without deserialization.
## Future Integration
This example demonstrates REFRAG concepts without modifying ruvector-core. For production use, consider:
1. **Phase 1**: Add `RefragEntry` as new struct in ruvector-core
2. **Phase 2**: Integrate policy network into ruvector-router
3. **Phase 3**: Update REST API with hybrid response format
See [Issue #10](https://github.com/ruvnet/ruvector/issues/10) for the full integration proposal.
## References
- [REFRAG: Rethinking RAG based Decoding (arXiv:2509.01092)](https://arxiv.org/abs/2509.01092)
- [RuVector Documentation](https://github.com/ruvnet/ruvector)