| --- |
| base_model: |
| - google/gemma-3-1b-pt |
| pipeline_tag: feature-extraction |
| --- |
| # Training Gemma-3-1B Embedding Model with LoRA |
|
|
| In our previous post, [Training a Query Fan-Out Model](https://dejan.ai/blog/training-a-query-fan-out-model/), we demonstrated how to generate millions of high-quality query reformulations without human labelling, by navigating the embedding space between a seed query and its target document and then decoding each intermediate vector back into text using a trained query decoder. |
|
|
| That decoder's success critically depends on having an embedding encoder whose latent geometry is fully under our control: off-the-shelf models (e.g. mxbai embed large) optimize for general semantic similarity, not for invertibility, so their embeddings cannot reliably be mapped back into meaningful queries. |
|
|
| To bridge that gap, this article introduces **Gemma-Embed**, a bespoke 256-dim embedding model built by fine-tuning `google/gemma-3-1b-pt` with LoRA adapters and contrastive objectives. By training our own encoder, we lock in a consistent, L2-normalized latent space that the subsequent query decoder can invert with high fidelity. |
|
|
| ## Quick Recap of the Query Fan-Out Mission |
|
|
| - **Query Decoder**: Train a T5-based model to invert a fixed retrieval encoder (e.g. GTR) so that any embedding vector produces the original query. Achieved ~96% cosine similarity on reconstruction. |
| - **Latent Space Traversal**: For each (query, document) pair, interpolate in the embedding space, decode each point, and retain reformulations that improve retrieval metrics—yielding hundreds of thousands of synthetic examples. |
| - **Production Model (qsT5)**: Fine-tune T5 on that synthetic dataset (with and without pseudo-relevance feedback) to internalize traversal patterns—so at inference time it generates diverse, effective reformulations without any explicit vector arithmetic. |
|
|
| Together, these steps automate query fan-out, boost retrieval performance, and open the door to interpretable, language-agnostic search suggestions. |
|
|
| To power a query fan‑out decoder that inverts embeddings back to natural language queries, we need an embedding encoder whose latent geometry we control. Since no off‑the‑shelf Gemma‑3 embedding model exists, we fine‑tune `google/gemma‑3‑1b‑pt` with LoRA and contrastive objectives to produce high‑quality, L2‑normalized 256‑dim embeddings. |
|
|
| ## Model Architecture |
|
|
| ### Base Encoder |
| - `google/gemma-3-1b-pt` (1 B params) |
|
|
| ### LoRA Adapters |
| - **Target modules**: `q_proj`, `v_proj` |
| - **Rank (r)**: 8 |
| - **Alpha (α)**: 16 |
| - **Dropout**: 0.05 |
|
|
| ### Projection Head |
| - **Input**: hidden_size (1024) |
| - **MLP**: Linear(1024→512) → ReLU → Linear(512→256) |
| - **L2 normalization** |
| |
| ## Data and Format |
| |
| ### Phase 1 – Unsupervised SimCSE |
| - **Source**: `text.txt` (wiki sentences or plain text logs) |
| - **Size**: 579,719 sentences |
| - **Format**: UTF‑8 plain text, one sequence per line |
| - **Sample lines**: |
| ``` |
| Breaking news: stock markets rally as central bank hints at rate cut. |
| How do I fine‑tune a large language model for embeddings? |
| The Northern Lights are visible tonight in high‑latitude regions. |
| ``` |
| |
| ### Phase 2 – Supervised Paraphrase Contrastive |
| - **Source**: `triplets.csv` |
| - **Columns**: `a_ids,a_mask,p_ids,p_mask,n_ids,n_mask` (token IDs & masks) |
| - **Size**: user‑provided paraphrase pairs (e.g., ParaNMT ~3.6 M, QuoraQP ~400 k, PAWS ~60 k) |
| - **Format**: CSV with header. Each row: |
| ``` |
| a_ids,a_mask,p_ids,p_mask,n_ids,n_mask |
| 102 345 ... ,1 1 ... ,203 456 ... ,1 1 ... ,307 523 ... ,1 1 ... |
| ``` |
| - Original text pairs stored in scripts folder for reference. |
| |
| ### Phase 3 – In‑Domain Self‑Contrast |
| - **Source**: `queries.db` |
| ```sql |
| CREATE TABLE queries ( |
| query_id INTEGER PRIMARY KEY AUTOINCREMENT, |
| query TEXT UNIQUE NOT NULL |
| ); |
| ``` |
| - **Size**: 7,129,444 unique queries |
| - **Pretokenized**: `pretokenized_queries.pt` |
| - **Tensors**: `input_ids` (7,129,444 × 128), `attention_mask` (7,129,444 × 128) |
| - **File size**: ~13.5 GB |
| - **Sample queries**: |
| ```sql |
| SELECT query FROM queries LIMIT 5; |
| How to bake sourdough at home? |
| Weather tomorrow in Sydney |
| Best restaurants near me open now |
| convert 1 mile to kilometers |
| streamlit file uploader example |
| ``` |
| |
| ## Training Pipeline |
|
|
| | Phase | Objective | Loss | Batch | Epochs | LR | Data Size | |
| |-------|-----------|------|-------|--------|----|-----------| |
| | 1 | Unsupervised SimCSE | InfoNCE (τ=0.05) | 12 | 1 | 1e‑5 | 579,719 sentences | |
| | 2 | Supervised Triplet Contrastive | TripletMarginLoss(0.2) | 12 | 3 | 1e‑5 | ~4 M triplets | |
| | 3 | In‑Domain Self‑Contrast | InfoNCE (τ=0.05) | 64 | 1 | 1e‑5 | 7,129,444 queries | |
|
|
| ## File Layout |
|
|
| ``` |
| train-gemma/ |
| ├── text.txt |
| ├── triplets.csv |
| ├── queries.db |
| ├── pretokenized_queries.pt |
| ├── scripts/ |
| │ ├── train_stage_1.py |
| │ ├── train_stage_2.py |
| │ ├── pretokenize_queries.py |
| │ └── train_stage_3.py |
| ├── stage1_simcse/final/ |
| ├── phase2_triplet_amp/final/ |
| └── phase3_self_contrast/final/ |
| ``` |
|
|
| ## Sample Data Sizes |
|
|
| - **text.txt**: 579,719 lines (~50 MB) |
| - **triplets.csv**: depends on sources (~500 MB for 4 M rows) |
| - **queries.db**: ~200 MB SQLite file |
| - **pretokenized_queries.pt**: 13.5 GB |
| |
| ## Inference Test |
| |
|  |
