deepx-embedding

Vietnamese Legal Document Retrieval — State-of-the-Art

🌐 Blog Post | 💻 GitHub

DeepX Embedding v1.0 is a 772M parameter embedding model optimized for Vietnamese legal document retrieval. It combines Gated DeltaNet-2 linear attention (O(n)) with Hyperloop weight sharing to achieve strong retrieval quality while maintaining constant throughput regardless of sequence length.

nDCG@10 = 0.8162 on Zalo Legal Text Retrieval — surpassing previous SOTA (0.7813) by +4.5%.

Benchmark Results

Key Features

• Linear attention O(n) — Gated DeltaNet-2: processes 8K tokens with same VRAM as 512 tokens
• Hyperloop architecture — 35 compute passes from only 9 unique layer parameter sets
• Matryoshka embeddings — Quality at any dimension from 256d to 1536d

• ColBERT dual output — Single vector (1536d) for ANN search + token vectors (128d) for MaxSim reranking
• Custom vocabulary — 186,046 tokens optimized for Vietnamese + English
• YaRN RoPE — 8K tokens validated, 128K supported

Architecture

Input text
  → Custom Tokenizer (186,046 vocab)
  → Frozen Token Embedding (186046 × 1536)
  → Begin Block: 4 unique NarrowA layers
  → Phase1 Loop ×2: [WideA + NarrowA×4] per iteration = 10 passes
  → Phase2 Loop ×4: [NarrowB×4 + WideB] per iteration = 20 passes
  → End Block: 1 unique WideB layer
  → RMSNorm → Attention Pooling → 1536-d vector

Total: 35 compute passes. Per-loop LoRA + RoDE (Rotary Depth Embedding) differentiate each iteration.

Model Size

Gated DeltaNet-2 (GDN-2)

Pure linear attention with O(n) complexity. Each layer maintains a running state updated via learned decay, erase, and write gates:

state_t = decay_t * state_{t-1}
state_t -= erase_t * (erase_t @ state_t - write_t * v_t)
output_t = q_t @ state_t

No KV cache, no quadratic slowdown. Uses FLA (flash-linear-attention) Triton kernels for efficient chunk-parallel training.

Training

Training pipeline: conservative training → long-sequence expose (4K-8K) → hard negative mining → domain boost.

Usage

import torch
from transformers import AutoTokenizer
from modeling.pipeline import DeepXPipeline
from config import DeepXConfig

# Load
tokenizer = AutoTokenizer.from_pretrained("dxtech-asia/deepx-embedding-v1")
config = DeepXConfig()
pipeline = DeepXPipeline.from_pretrained(config, "deepx_v1.0.pt")
pipeline.eval().cuda()

# Encode
text = "Mức phạt khi vượt đèn đỏ là bao nhiêu?"
inputs = tokenizer(text, return_tensors="pt", max_length=8192, truncation=True)
id_remap = torch.load("id_remap.pt")
input_ids = id_remap[inputs["input_ids"]]

with torch.no_grad():
    embedding = pipeline.encode(input_ids.cuda(), inputs["attention_mask"].cuda())
    # Shape: (1, 1536), L2-normalized

Inference Speed

On RTX 5070 Ti, FP16 inference.

Citation

@misc{deepx2026embedding,
  title={DeepX Embedding v1.0: Vietnamese Legal Retrieval with Gated DeltaNet-2 Linear Attention},
  author={DX Tech Asia},
  year={2026},
  url={https://huggingface.co/dxtech-asia/deepx-embedding-v1}
}

License

Apache 2.0