WRAG 2.0: Weight-Retrieval Augmented Generation

WRAG 2.0 is a novel neural architecture that uses dynamic weight retrieval during the forward pass to achieve domain specialization without retraining the base model.

🎯 Performance

Trained on 3 domains with only 500 samples per domain (5 epochs):

Domain	Accuracy
Medical (PubMedQA)	70%
Legal (LexGLUE)	84%
Code (CodeSearchNet)	95%
Average	83%

🚀 Quick Start

from modeling_wrag2 import WRAG2TextModel
import torch

# Load model
model = WRAG2TextModel(
    base_model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    num_shards=10,
    k=3,
    num_wr_layers=3
)

# Load trained weights
state_dict = torch.load("pytorch_model.bin")
model.load_state_dict(state_dict, strict=False)

# Move to device
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()

# Inference
text = ["Question: What is the treatment for diabetes?"]
inputs = model.tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=256).to(device)

with torch.no_grad():
    logits, shard_scores = model(**inputs)
    prediction = logits.argmax(dim=-1)
    
print(f"Prediction: {prediction.item()}")  # 0 or 1

💡 Key Features

✅ Domain Specialization: Different weight shards activate for different domains
✅ Efficient: Only 264M trainable params (20% of total model)
✅ Fast Training: 5 epochs, ~30 minutes on L4 GPU
✅ Memory Efficient: Fits in 24GB GPU
✅ Lightweight Download: Only 264MB (base model downloaded separately)

🏗️ Architecture

Input Text
    ↓
Frozen TinyLlama-1.1B (feature extraction)
    ↓
Mean Pooling
    ↓
Weight Retrieval Layer 1 (10 shards, k=3)
    ↓ ReLU
Weight Retrieval Layer 2 (10 shards, k=3)
    ↓ ReLU
Weight Retrieval Layer 3 (10 shards, k=3)
    ↓ ReLU
Classification Head (2 classes)
    ↓
Output

📊 Model Details

Base Model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 (frozen)
Trainable Parameters: 264M (20.2% of total)
Total Parameters: 1.3B
Weight Shards: 10 per layer
Top-k Selection: 3 shards per forward pass
Number of WR Layers: 3

🔧 Training Details

Optimizer: AdamW
Learning Rate: 1e-4
Batch Size: 4
Epochs: 5
Max Sequence Length: 256
Training Data: 500 samples per domain (medical, legal, code)
Training Time: ~30 minutes on NVIDIA L4 GPU

📦 Installation

pip install torch transformers huggingface_hub

🎓 How It Works

WRAG 2.0 uses dynamic weight retrieval:

Input text is encoded by frozen TinyLlama
Each weight retrieval layer:
- Computes similarity between input and 10 weight shards
- Selects top-3 most relevant shards
- Composes dynamic weight matrix via weighted combination
- Applies dynamic weights to input
Classification head produces final prediction

This allows the model to specialize for different domains without retraining the base model!

⚠️ Limitations

Binary classification only (yes/no, 0/1)
Trained on limited data (500 samples per domain)
Requires TinyLlama base model (~4.4GB) to be downloaded separately
Best for domain-specific tasks similar to training domains

📝 Citation

@misc{wrag2-2025,
  title={WRAG 2.0: Weight-Retrieval Augmented Generation},
  year={2025},
  url={https://huggingface.co/YOUR_USERNAME/wrag2-text-classifier}
}

📄 License

Apache 2.0

🤝 Contributing

Issues and pull requests welcome!

Downloads last month: 6

Gyeti123
/

wrag2-text-classifier