mz-RAG-v9-squad

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

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

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.