You can deploy this model and start earning money today!

honey0119/
mz-RAG-v9-squad

Views

No views yet

mz-RAG v9 SQuAD Experiment

Query-to-z alignment experiment using SQuAD 1.1 dataset.

Experiment Overview

Goal: Learn z embeddings for document reconstruction, then align query embeddings to z space for retrieval.

Phase 1: z Reconstruction

Train z embeddings to reconstruct SQuAD paragraphs via NLL loss
LLM (Llama-3.2-3B) is frozen, only z embeddings are trained

Phase 2: Query Encoder Alignment

Freeze z embeddings from Phase 1
Train a projection head (MLP) to map query embeddings to z space
Use InfoNCE contrastive loss

Results

Metric	Score
Top-1 Accuracy	32.5%
Top-5 Accuracy	56.0%
Top-10 Accuracy	63.5%
MRR	0.433

Test set: 200 held-out queries (1 per document)

Configuration

Model: meta-llama/Llama-3.2-3B (4bit quantized)
Documents: 200 SQuAD paragraphs (min 4 questions each)
Max doc length: 128 tokens
Query split: 3 train / 1 test per document
Projection head: 512 hidden dim MLP
Training: 5 epochs, lr=1e-3, batch_size=16

Files

z_checkpoint/: Phase 1 z embeddings checkpoint
query_encoder/: Phase 2 query encoder checkpoint
results/: Evaluation results and examples

Usage

python

from huggingface_hub import hf_hub_download

# Download z checkpoint
z_path = hf_hub_download("honey0119/mz-RAG-v9-squad", "z_checkpoint/checkpoint_final.pt")

# Download query encoder
encoder_path = hf_hub_download("honey0119/mz-RAG-v9-squad", "query_encoder/query_encoder.pt")

Citation

Part of the mz-RAG project for document retrieval via learned z embeddings.