| ---
|
| language:
|
| - vi
|
| - en
|
| license: apache-2.0
|
| tags:
|
| - embedding
|
| - vietnamese
|
| - legal
|
| - retrieval
|
| - linear-attention
|
| - gdn2
|
| - matryoshka
|
| library_name: pytorch
|
| pipeline_tag: feature-extraction
|
| model-index:
|
| - name: deepx-embedding-v09
|
| results:
|
| - task:
|
| type: Retrieval
|
| dataset:
|
| name: Zalo Legal Text Retrieval
|
| type: GreenNode/zalo-ai-legal-text-retrieval-vn
|
| metrics:
|
| - type: ndcg_at_10
|
| value: 0.7449
|
| - type: mrr_at_10
|
| value: 0.6921
|
| - type: recall_at_10
|
| value: 0.9086
|
| ---
|
|
|
| # DeepX Embedding 0.9 (Preview)
|
|
|
| ## Model Description
|
|
|
| > ⚠️ **Preview Release** — This is a preview version for testing and evaluation. Not recommended for production use. Final v1.0 release will include improved quality and optimizations.
|
|
|
| DeepX Embedding 0.9 is a Vietnamese-focused embedding model built on a novel **Gated DeltaNet-2 (GDN-2)** linear attention architecture with **O(n) complexity**. It uses Hyperloop weight sharing (9 unique layers, 35 compute passes) with architecture design inspired by Google Gemma 4 E2B.
|
|
|
| **Key features:**
|
| - O(n) linear attention via FLA (flash-linear-attention) Triton kernels
|
| - Matryoshka embedding: supports 256, 512, 768, 1024, 1536 dimensions
|
| - Trained for Vietnamese legal document retrieval
|
| - Tokenizer and token embedding from Gemma 4 E2B (frozen, not trained)
|
| - Effective model size: 486M trainable parameters (token embedding excluded)
|
| - Finetune and quantization only apply to 486M backbone parameters
|
| - Competitive with 560M-600M parameter SOTA models
|
|
|
| ## Architecture
|
|
|
| | Component | Detail |
|
| |---|---|
|
| | Base | Gemma 4 E2B (262K vocab, 1536 hidden) |
|
| | Attention | GDN-2 (Gated DeltaNet-2), pure linear O(n) |
|
| | Structure | Hyperloop: Begin(4) + Phase1×2(5) + Phase2×4(5) + End(1) = 35 passes |
|
| | Unique layers | 9 (shared via loop + per-iteration LoRA) |
|
| | Total params | 889M (486M trainable backbone + 403M frozen token embedding) |
|
| | Embedding dim | 1536 (Matryoshka: 256/512/768/1024/1536) |
|
| | Max sequence | 2048 tokens (training), 131K (theoretical via YaRN RoPE) |
|
| | Pooling | Attention-weighted pooling |
|
| | Depth signal | RoDE (Rotary Depth Encoding) |
|
|
|
| ### GDN-2 Attention
|
|
|
| Unlike standard softmax attention (O(n²)), GDN-2 uses a gated delta rule recurrence computed via chunked parallel scan. This gives:
|
| - **O(n) time and memory** for any sequence length
|
| - Constant memory per token (no KV cache growth)
|
| - Hardware-efficient via FLA Triton kernels
|
|
|
| ### Hyperloop Weight Sharing
|
|
|
| The model reuses 9 unique layer parameter sets across 35 compute passes:
|
| - 4 NarrowA layers (8 heads, MLP 6144)
|
| - 4 NarrowB layers (8 heads, MLP 12288)
|
| - 1 WideA layer (16 heads, MLP 6144)
|
| - 4 WideB layers (16 heads, MLP 12288)
|
|
|
| Per-iteration LoRA (rank 16) differentiates each pass within a loop.
|
|
|
| ## Training
|
|
|
| ### Data
|
| - Vietnamese legal query-passage pairs (500K+)
|
| - Vietnamese legal news retrieval pairs
|
| - English scientific retrieval pairs
|
| - Custom Vietnamese domain-specific data
|
| - Custom hard negative mining from Vietnamese legal corpus (61K documents)
|
| - Iterative hard negative refinement (multiple rounds, rank 1-50)
|
|
|
| ### Training Details
|
| - Optimizer: AdamW 8-bit
|
| - Loss: InfoNCE + Matryoshka (multi-dim)
|
| - Trainable params: 486M / 889M (54.7%)
|
| - Gradient checkpointing enabled
|
| - Pure GDN-2 (alpha=0, softmax skipped)
|
|
|
| ## Evaluation
|
|
|
| ### Zalo Legal Text Retrieval
|
|
|
| | Metric | DeepX v0.9 | mE5-large | vietlegal-harrier (SOTA) |
|
| |---|---|---|---|
|
| | **nDCG@10** | **0.7449** | 0.6660 | 0.7813 |
|
| | MRR@10 | 0.6921 | — | 0.7303 |
|
| | Recall@10 | 0.9086 | — | 0.9321 |
|
|
|
| ### Speed (O(n) advantage)
|
|
|
| | Sequence Length | DeepX (GDN-2) | Softmax Equivalent |
|
| |---|---|---|
|
| | 512 | 0.31 it/s | 0.31 it/s |
|
| | 2048 | 0.18 it/s | 0.01 it/s |
|
| | 8192 | ~0.18 it/s | OOM |
|
|
|
| Linear attention maintains constant speed regardless of sequence length.
|
|
|
| ## Usage
|
|
|
| ```python
|
| import torch
|
| from transformers import AutoTokenizer
|
|
|
| # Load
|
| tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-E2B-it")
|
| # Load model (see repo for full pipeline code)
|
| model = load_deepx_model("deepx_v09.pt")
|
|
|
| # Encode
|
| text = "Mức phạt khi vượt đèn đỏ là bao nhiêu?"
|
| inputs = tokenizer(text, return_tensors="pt", max_length=2048, truncation=True)
|
| with torch.no_grad():
|
| embedding = model(inputs["input_ids"], attention_mask=inputs["attention_mask"], normalize=True)
|
| # embedding.shape = (1, 1536)
|
|
|
| # Matryoshka: use first N dims
|
| embedding_256d = embedding[:, :256] # 90% quality, 6x less storage
|
| ```
|
|
|
| ## Intended Use
|
|
|
| - Vietnamese legal document retrieval
|
| - Vietnamese question-answering (retrieval component)
|
| - Cross-lingual retrieval (Vietnamese ↔ English)
|
| - General Vietnamese semantic search
|
|
|
| ## Limitations
|
|
|
| - Primarily trained on Vietnamese legal domain; general-domain performance may be lower
|
| - Requires FLA library (Triton kernels) for efficient inference
|
| - Not trained on conversational/chat data
|
| - English performance limited (secondary training only)
|
|
|
| ## Hardware Requirements
|
|
|
| - **Inference**: RTX 3060 12GB+ (model ~3.5GB VRAM)
|
| - **Training**: RTX 5070 Ti 16GB+ recommended
|
| - **Dependencies**: PyTorch 2.0+, transformers, fla (flash-linear-attention), triton
|
|
|
| ## Citation
|
|
|
| ```bibtex
|
| @misc{deepx2026,
|
| title={DeepX: Vietnamese Embedding Model with Gated DeltaNet-2 Linear Attention},
|
| author={DXTech Asia},
|
| year={2026},
|
| url={https://huggingface.co/dxtech-asia/deepx-embedding-v09}
|
| }
|
| ```
|
|
|
| ## License
|
|
|
| Apache 2.0 (code) / Model weights follow Gemma license terms.
|
|
|
