Technical Report: DeepX Embedding v1.0

1. Introduction

DeepX Embedding v1.0 is a Vietnamese embedding model designed for legal document retrieval. It combines a novel linear attention mechanism (Gated DeltaNet-2) with Hyperloop weight sharing to achieve strong retrieval quality while maintaining O(n) time complexity for long sequences.

The model was trained on 2x RTX 5070 Ti 16GB GPUs over roughly 600 hours. It achieves state-of-the-art on the Zalo Legal Text Retrieval benchmark, beating models that had access to much larger compute budgets.

This report documents the architecture decisions, training procedure, and evaluation results for reproducibility and community reference.

2. Architecture Overview

DeepX v1.0 uses a custom tokenizer with 186,046 tokens optimized for Vietnamese and English. The vocabulary was independently constructed to maximize coverage for Vietnamese legal text and English technical content, while minimizing unused tokens from other language families. Representations pass through a custom backbone with 35 compute passes organized via Hyperloop weight sharing.

2.1 High-level Structure

Input text
  → Custom Tokenizer (186,046 vocab, optimized for VN + EN)
  → Frozen Token Embedding (186046 x 1536)
  → Begin Block: 4 unique NarrowA layers
  → Loop 1 x2: [WideA + NarrowA x4] per iteration
  → Loop 2 x4: [NarrowB x4 + WideB] per iteration
  → End Block: 1 unique WideB layer
  → RMSNorm
  → Attention Pooling → single vector (1536-d)
  → ColBERT Head → token vectors (T x 128-d)

2.2 Model Size

The token embedding is frozen throughout training. The vocabulary was constructed with 186,046 tokens covering Vietnamese, English, code symbols, and common Unicode ranges, while excluding CJK, Hangul, Arabic, Thai, and other scripts unused in the target domains. This reduces embedding size by 29% compared to a full multilingual tokenizer with zero impact on VN/EN quality.

2.3 Layer Types

3. Gated DeltaNet-2 (GDN-2) Attention

The core attention mechanism is GDN-2, a pure linear attention variant with O(n) complexity based on the gated delta rule recurrence.

3.1 How GDN-2 Works

Each attention layer maintains a running key-value state that gets updated at every token via three learned gates:

• Decay gate: Controls how much of the existing state to retain
• Erase gate: Selectively removes outdated information from the state
• Write gate: Controls what new information gets written into the state

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

A short 1D causal convolution on Q and K provides local context before the recurrence. An output gate (SiLU) is applied after for additional expressiveness.

3.2 Why Linear Attention Matters

3.3 FLA Kernel Integration

GDN-2 uses the fla library (flash-linear-attention) Triton kernels for chunk-parallel computation during training. When FLA is not available, the model falls back to a sequential implementation that is functionally identical but slower.

4. Hyperloop Weight Sharing

Our architecture is based on the Hyperloop Transformer [17], which organizes a looped transformer into begin, middle, and end blocks where only the middle block is applied recurrently.

4.1 Structure

• Begin: 4 unique NarrowA layers
• Loop 1: 2 iterations x [1 WideA + 4 NarrowA] = 10 passes
• Loop 2: 4 iterations x [4 NarrowB + 1 WideB] = 20 passes
• End: 1 unique WideB layer

9 unique parameter sets producing 35 compute passes.

4.2 Per-Loop LoRA (DX contribution)

Each iteration receives its own low-rank adapter (LoRA, rank 16) applied to all projection matrices (Q, K, V, O, gate, up, down). This differentiates entire multi-layer blocks (5 layers per iteration) simultaneously, preventing depth collapse in looped architectures.

4.3 RoDE (Rotary Depth Embedding, DX contribution)

We introduce RoDE, a novel positional signal that encodes loop iteration depth directly into Q and K via rotation. Standard RoPE only encodes token position. RoDE adds a second rotation based on loop index, giving the attention mechanism an explicit depth signal. Uses 16 dimensions (8 rotary pairs). To our knowledge this has not been proposed before.

5. Training

5.1 Hardware and Optimization

Pipeline parallel splits the backbone: GPU0 handles begin + loop 1 (14 passes), GPU1 handles loop 2 + end (21 passes).

5.2 Iterative Training

Training followed an iterative approach over 14-15 rounds, each round building on the previous checkpoint through cycles of training, evaluating, mining new hard negatives, and adjusting data mixture. Total training time: around 600 hours.

5.3 Training Data

Cross-lingual (EN-VI):

Vietnamese General:

Legal Domain:

5.4 Loss Function

• InfoNCE loss with temperature 0.07
• Matryoshka loss at dims 256, 512, 768, 1024, 1536
• ColBERT MaxSim loss on token vectors

6. Evaluation Results

6.1 Zalo Legal Text Retrieval

Dataset: GreenNode/zalo-ai-legal-text-retrieval-vn (61,396 docs, 818 test queries)

DeepX v1.0 surpasses the previous SOTA by +3.5 points on nDCG@10.

6.2 Matryoshka Embedding Quality

256-d embeddings provide 6x storage reduction with minimal quality loss.

7. Outputs and Usage

7.1 Dual Output

1. Single vector (1536-d): For ANN search. L2-normalized, cosine similarity.
2. ColBERT token vectors (T x 128-d): For MaxSim reranking. Per-token L2-normalized.

7.2 Two-Stage Retrieval

Stage 1: Query vector vs. corpus vectors → ANN (FAISS/Milvus) → top-100
Stage 2: Query tokens vs. top-100 doc tokens → MaxSim → top-10

7.3 Code Example

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

tokenizer = AutoTokenizer.from_pretrained("dxtech-asia/deepx-embedding-v1")
config = DeepXConfig()
pipeline = DeepXPipeline.from_pretrained(
    config,
    backbone_path="checkpoints/deepx_v1.pt",
    embed_path="deploy/pruned/token_embedding_pruned.pt",
)
pipeline.eval()

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("deploy/pruned/tokenizer/id_remap.pt")
input_ids = id_remap[inputs["input_ids"]]

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

8. Intended Use and Limitations

8.1 Intended Use

• Vietnamese legal document retrieval and QA
• Vietnamese semantic search
• Cross-lingual retrieval (Vietnamese and English)
• General Vietnamese text similarity

8.2 Limitations

• Primarily trained on Vietnamese legal text; other domains not thoroughly tested
• Requires fla library for efficient inference; falls back to slower sequential mode without it
• Not trained on conversational/chat data
• English quality is secondary

8.3 Hardware Requirements

• Inference: Recommended GPU with 4GB+ VRAM. Model weights ~1.5GB in FP16. Can be quantized to FP8 to reduce further.
• Training: 2x RTX 5070 Ti 16GB (pipeline parallel)
• Dependencies: PyTorch 2.0+, transformers, fla, triton, einops

9. Citation

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

10. License

Code: Apache 2.0 | Model weights: Apache 2.0 | Tokenizer: Custom vocabulary.

11. References

1. S. Yang, et al. "Gated DeltaNet." 2024. arXiv:2412.06464.
2. S. Yang, et al. "Decoupling Erase and Write in Linear Attention (GDN-2)." 2025. arXiv:2605.22791.
3. S. Yang, et al. "Parallelizing Linear Transformers with the Delta Rule." ICML 2025. arXiv:2406.06484.
4. S. Yang, B. Wang, et al. "Gated Linear Attention Transformers." ICML 2024. arXiv:2312.06635.
5. S. Yang, et al. "Flash Linear Attention (FLA)." GitHub.
6. E. J. Hu, et al. "LoRA: Low-Rank Adaptation of Large Language Models." ICLR 2022.
7. O. Khattab, M. Zaharia. "ColBERT." SIGIR 2020.
8. L. Wang, et al. "Text Embeddings by Weakly-Supervised Contrastive Pre-training (E5)." ACL 2024.
9. A. Kusupati, et al. "Matryoshka Representation Learning." NeurIPS 2022.
10. B. Peng, et al. "YaRN: Efficient Context Window Extension." ICLR 2024.
11. GreenNode. "Zalo AI Legal Text Retrieval VN." HuggingFace.
12. mainguyen9. "vietlegal-harrier-0.6b." HuggingFace.
13. phamson02. "large-vi-legal-queries." HuggingFace.
14. T. Dettmers, et al. "8-bit Optimizers via Block-wise Quantization." ICLR 2022.
15. N. Muennighoff, et al. "MTEB: Massive Text Embedding Benchmark." EACL 2023.
16. A. Zeitoun, et al. "Hyperloop Transformers." 2026. arXiv:2604.21254.
17. S. Davis, et al. "Effective Parameter Sharing with Layer-wise LoRA." Google DeepMind, 2024. arXiv:2410.20672.

Acknowledgments

This work uses the FLA library from the flash-linear-attention team, and the Zalo Legal Text Retrieval benchmark from GreenNode. We acknowledge the vietlegal-harrier team for establishing a strong baseline that motivated our work on hard negative mining.