---
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.