Cola-DLM

English · 中文

Cola DLM (Continuous Latent Diffusion Language Model) is a hierarchical continuous latent-space diffusion language model. It combines a Text VAE with a block-causal Diffusion Transformer (DiT) prior: the VAE maps text into continuous latent sequences and decodes latents back to tokens, while the DiT performs latent prior transport through Flow Matching.

This model repository contains the HuggingFace-format checkpoint for the paper Continuous Latent Diffusion Language Model.

Links

• Model repository:
• GitHub repository:
• Paper:
• HuggingFace Daily Paper:
• Project page:
• Blog post:
• Zhihu article:

Model Files

The expected repository layout is:

.
├── cola_dlm/
│   ├── cola_dit/
│   │   ├── config.json
│   │   └── model.safetensors*
│   └── cola_vae/
│       ├── config.json
│       └── model.safetensors*
├── tokenizer.json
├── README.md
└── README_zh.md

The checkpoint consists of two cooperating modules:

• ColaDiTModel: a block-causal 1-D Diffusion Transformer prior over continuous text latents.
• ColaTextVAEModel: a Text VAE encoder and conditional decoder for text-to-latent and latent-to-text mapping.

Quickstart

Install the Cola DLM code package from the GitHub repository, then install the download helper:

git clone https://github.com/ByteDance-Seed/Cola-DLM.git
cd Cola-DLM
pip install -e .
pip install huggingface_hub

Download the model files:

huggingface-cli download ByteDance-Seed/Cola-DLM --local-dir hf_models

Run a minimal Python example:

import torch
from tokenizers import Tokenizer

from cola_dlm import (
    ColaDiTModel,
    ColaTextVAEModel,
    generate_task_repaint_inference,
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

dit = ColaDiTModel.from_pretrained("hf_models/cola_dlm/cola_dit").to(device)
vae = ColaTextVAEModel.from_pretrained("hf_models/cola_dlm/cola_vae").to(device)
tokenizer = Tokenizer.from_file("hf_models/tokenizer.json")

prompts = [{"question": "Question: What is the capital of France? Answer:"}]
results = generate_task_repaint_inference(
    dit=dit,
    vae=vae,
    tokenizer=tokenizer,
    prompts=prompts,
    task_name="lambada",
    device=device,
    max_new_tokens=32,
    temperature=0.0,
    guidance_scale=7.0,
    timestep_num=16,
    pad_token_id=100277,
)

print(results[0]["generate"])

OpenAI-Compatible Serving

The companion openai_adapter/ service in the Cola DLM code release exposes this model through an OpenAI-compatible Chat Completions endpoint:

POST /v1/chat/completions

Install the adapter dependencies from the code repository root:

pip install -e .
pip install -r openai_adapter/requirements.txt

Start the service:

export COLA_DIT_PATH=hf_models/cola_dlm/cola_dit
export COLA_VAE_PATH=hf_models/cola_dlm/cola_vae
export COLA_TOKENIZER_PATH=hf_models/tokenizer.json
export COLA_MODEL_NAME=cola-dlm
export COLA_API_KEY=change-me

uvicorn openai_adapter.server:app --host 0.0.0.0 --port 8000

Then send a request:

curl http://127.0.0.1:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer change-me" \
  -d '{
    "model": "cola-dlm",
    "messages": [
      {
        "role": "user",
        "content": "Question: What is the capital of France? Answer:"
      }
    ],
    "temperature": 0,
    "max_tokens": 32,
    "stream": false
  }'

The adapter currently supports non-streaming completions.

Model Details

• Architecture: Text VAE + block-causal DiT latent prior.
• Training objective: two-stage training with Text VAE pretraining followed by joint Text VAE + DiT training using Flow Matching.
• Training-compute checkpoint: the released weights correspond to the 2000 EFLOPs checkpoint reported in the paper's RQ4 scaling curve.
• Tokenizer: OLMo 2 tokenizer with a 100,278-entry vocabulary.
• Special token ids: pad_token_id=100277, eos_token_id=100257, im_end_token_id=100265.
• Framework: PyTorch 2.1+ and HuggingFace Transformers 4.40+.
• License: Apache License 2.0.

Evaluation

Reference zero-shot benchmark results from the open-source inference implementation:

The open-source HuggingFace Transformers implementation may differ slightly from the internal implementation used in the paper, so per-task numbers can fluctuate slightly. The overall trend is consistent with the paper.

Intended Use

Cola DLM is intended primarily for research on hierarchical latent-variable language models, continuous latent diffusion for text, Flow Matching priors, and benchmark-style text generation.

This checkpoint is not instruction-tuned and has not gone through RLHF. It should not be treated as a production chatbot or used for safety-critical decision making.

Limitations

• The model was trained primarily on English text; other languages are not well evaluated.
• Outputs may contain factual errors, offensive content, bias, or hallucinations.
• Generation quality can be sensitive to prompt format and prompt length. QA-style prompts such as "Question: ... Answer:" are recommended for quick evaluation.
• The model uses mutable KV caches during generation; service implementations should serialize generation inside one process unless cache handling is explicitly isolated.

Safety Statement and Use Restrictions

Cola DLM is a research-oriented checkpoint for continuous latent diffusion language modeling. The released model is relatively small and has not been instruction-tuned, RLHF-aligned, or systematically safety-aligned. Therefore, it does not provide reli