Embedding Inversion via Conditional Masked Diffusion  

See also: Embedding Fingerprints -- identify which embedding model produced a vector using digit-level tokenization.

Text embeddings are widely assumed to be safe, irreversible representations. This project demonstrates otherwise: given only an embedding vector, we reconstruct the original text using conditional masked diffusion.

How It Works

Existing inversion methods (Vec2Text, ALGEN, Zero2Text) generate tokens autoregressively and require iterative re-embedding through the target encoder. This creates two bottlenecks: attack cost scales with correction iterations, and left-to-right generation accumulates errors with no mechanism to revise earlier tokens.

We take a different approach: embedding inversion as conditional masked diffusion. Starting from a fully masked sequence, a denoising model reveals tokens at all positions in parallel, conditioned on the target embedding via adaptive layer normalization (AdaLN-Zero). Each denoising step refines all positions simultaneously using global context, without ever re-embedding the current hypothesis.

The approach is encoder-agnostic by construction. The embedding vector enters only through AdaLN modulation of layer normalization parameters, so the same architecture applies to any embedding model without alignment training or architecture-specific modifications.

Check out the live demo to see it in action.

Data Preparation

Two-stage pipeline: download raw texts, then tokenize and encode with the target embedding model.

# Stage 1: download multilingual raw texts (run once, reusable across encoders)
python prepare_data_fast.py --stage 1 \
  --langs en,zh,de,fr,ja,es,ru,ko,ar,pt \
  --n-samples 5000000

# Stage 2: tokenize + encode for a specific encoder (run per encoder)
python prepare_data_fast.py --stage 2 \
  --config configs/v2_qwen3.yaml \
  --encode-batch 512

Stage 1 downloads from C4/mC4 and saves raw texts as JSON chunks. Stage 2 reads those texts, tokenizes with the encoder's tokenizer (truncated to seq_len tokens), encodes into embeddings, and saves as numpy arrays.

Supported encoders via config:

Training

# Train with a specific config
python train.py --config configs/v2_qwen3.yaml

# Resume from checkpoint
python train.py --config configs/v2_qwen3.yaml --resume

Four encoders trained in parallel on A100-40GB GPUs. Qwen3-Embedding leads with 81.3% token accuracy at 72.5K steps. All models use the same 78M-parameter MDLM backbone with AdaLN-Zero conditioning.

Key training parameters (in config yaml):

• batch_size: per-GPU batch size (auto-detected if not set)
• grad_accum: gradient accumulation steps
• max_steps: total training steps
• lr: peak learning rate (cosine schedule with warmup)
• eval_every: validation frequency

The model uses mixed precision (bf16) and EMA weight averaging. Checkpoints are saved to checkpoints/ with best.pt tracking the lowest validation loss.

Training logs report step, loss, accuracy, and learning rate. Accuracy measures exact token match at masked positions.

Inference

Interactive Demo

pip install fastapi uvicorn
python demo_server.py

Opens on port 8080. Live demo: https://embedding-inversion-demo.jina.ai

Programmatic Inversion

python invert.py \
  --config configs/v2_qwen3.yaml \
  --checkpoint checkpoints/qwen3_best.pt \
  --text "The quick brown fox jumps over the lazy dog"

This encodes the input text, then runs diffusion decoding to reconstruct it from the embedding alone. Supports multiple decoding strategies (greedy, Euler sampling, adaptive re-masking).

Files

Citation

@article{xiao2026embedding,
  title={Embedding Inversion via Conditional Masked Diffusion Language Models},
  author={Xiao, Han},
  journal={arXiv preprint arXiv:2602.11047},
  year={2026}
}