Model Card: Cola DLM

July 20, 2026 · View on GitHub

English · 中文

This model card follows the structure suggested by Mitchell et al. (2019) Model Cards for Model Reporting. It describes the reference cola-dlm checkpoint released alongside this repository and the paper Continuous Latent Diffusion Language Model.

Paper

Model details

  • Model name: Cola DLM (Continuous Latent Diffusion Language Model)
  • Version: 0.1.0
  • Release date: 2026
  • Method: Hierarchical continuous latent diffusion language model — a Text VAE encoder q_phi(z_0 | x) + conditional decoder p_theta(x | z_0) paired with a block-causal Diffusion Transformer prior p_psi(z_0) learned by Flow Matching. Diffusion is used as latent prior transport (Eq. 2.1.4 of the paper), not token-level observation recovery.
  • Architecture: Two cooperating modules
    • ColaTextVAEModel — ~500M parameter Text VAE realizing q_phi and p_theta (4 encoder + 4 decoder blocks, dim=1536, ffn_dim=6144, latent_dim=16).
    • ColaDiTModel — ~1.8B parameter 1-D Diffusion Transformer realizing the block-causal latent prior p_psi (24 layers, txt_dim=emb_dim=2048, 16 heads × 128 head_dim, expand_ratio=4).
  • Training-compute checkpoint: This release ships the 2000 EFLOPs entry on the paper's RQ4 scaling curve — the largest-compute checkpoint in the paper. The total of ~2B parameters is strictly matched against the autoregressive and LLaDA baselines used in the same scaling comparison.
  • License: Apache License 2.0.
  • Framework: PyTorch 2.1+; HuggingFace Transformers 4.40+.
  • Precision: Master weights kept in fp32; forward passes use torch.autocast(dtype=torch.bfloat16) on CUDA.
  • Tokenizer: OLMo 2 tokenizer (BPE over a 100,278-entry vocabulary); pad_token_id=100277, eos_token_id=100257, im_end_token_id=100265.

Intended use

Primary use cases

  • Research on hierarchical / continuous latent text diffusion language models.
  • Closed-domain question answering and zero-/few-shot benchmark evaluation (e.g. LAMBADA, MMLU, HellaSwag, OBQA, RACE, SIQA, SQuAD, Story Cloze) — exactly the 8 tasks reported in the paper's RQ4.
  • Studying properties of latent-prior-transport text generation (controllability via CFG, block-wise re-painting, latent-space semantic structure, etc.).

Out-of-scope uses

  • Safety-critical decision-making (medical, legal, financial advice).
  • Generation of content that violates applicable laws, platform policies, or research-ethics norms.
  • Use as a drop-in chatbot without task-specific fine-tuning — the pretraining objective is hierarchical latent prior modeling, not instruction-following dialogue.

Factors

  • The model was pretrained primarily on English text; performance on other languages is untested and expected to be weak.
  • Generation quality depends strongly on prompt length relative to block_size (typically 16): prompts shorter than one block auto-disable classifier-free guidance for their first block and may produce weaker outputs. See docs/inference.md for the mitigation.

Metrics

Accuracy on 8 zero-shot benchmarks, batch size 70, timestep_num=16, guidance_scale=7.0, greedy decoding (temperature=0.0), max_new_tokens=32, 1000 samples per task, checkpoint global_step_300000 (the 2000 EFLOPs entry on the paper's RQ4 scaling curve):

TaskAccuracy (%)
LAMBADA50.80
MMLU19.30
OBQA23.00
HellaSwag10.70
RACE19.30
SIQA28.40
SQuAD30.90
Story Cloze30.77
Tasks Average26.65

Note on open-source model and accuracy: The released model weights correspond to the 2000 EFLOPs entry on the paper's RQ4 scaling curve — the largest training-compute checkpoint reported. Because the internal architecture used for evaluation in the paper differs slightly from the open-source HuggingFace Transformers-based implementation in this repository, per-task accuracy numbers may exhibit minor fluctuations, but the overall trend is consistent with the paper. Notably, the Tasks Average (26.65%) measured here is slightly higher than the final average reported in the paper.

See eval_output/accuracy_summary.csv and scripts/run_benchmark.sh for the exact protocol.

Training data

  • Large-scale public web and book text, filtered for English, deduplicated and quality-filtered. The corpus composition and mix ratios are documented in the paper; they are not included in this repository.
  • The training corpus is not redistributed as part of this repo; only the model weights and inference code are released.

Training procedure (overview)

Cola DLM is trained in two stages, summarized here for reference (training code is not in this repo; the open-source release is inference-only):

  1. Stage 1 — Text VAE pretraining (L_VAE in Eq. 2.2.1): reconstruction + BERT-style masking + KL regularization to a base prior, establishing a stable text↔latent correspondence.
  2. Stage 2 — Joint Text VAE + block-causal Text DiT pretraining (L_stage2 in Eq. 2.2.3): conditional Flow Matching loss L_FM for the block-causal prior p_psi, plus a regularized autoencoding term and a reference-encoder KL that suppresses latent drift.

See docs/architecture.md §9 for the explicit objectives.

Ethical considerations

  • Bias. Like any large-scale language model trained on internet text, Cola DLM will reflect and can amplify societal biases present in the training data (stereotypes, gendered assumptions, under-representation of minority dialects, etc.). Downstream users must audit outputs before deploying in user-facing products.
  • Toxicity & hallucinations. The model may produce factually incorrect, offensive, or harmful text, especially under adversarial prompting. Do not treat its outputs as authoritative.
  • Privacy. Training data was not de-duplicated at the document level beyond standard hashing; memorisation of rare text is possible.
  • Misuse. The model could be misused for generating spam, misinformation, phishing content, or non-consensual text about real people. Users must comply with local laws and platform terms of service.

Caveats & recommendations

  • Cola DLM is a research artifact. It ships without an instruction-tuning / RLHF stage, and its outputs will often underperform production chatbots.
  • For best quality, wrap short prompts in QA-style templates ("Question: ... Answer:") and use at least block_size tokens.
  • When deploying at scale, monitor output for safety / bias / factuality using independent classifiers, and surface disclaimers to end users.

Citation

See the README for the BibTeX entries (paper + open-source software).

Contact

Open an issue at https://github.com/your-org/cola-dlm/issues.