dxtech-asia/deepx-embedding-v09

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

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

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