amazon-science/ALoDLM

amazon-science/ALoDLM

Description: Research code and optimized inference for ALoDLM: Adaptively Looped Diffusion Language Models.

Language: Python

License: NOASSERTION

Stars: 2

Forks: 1

Open issues: 3

Created: 2026-10-02T23:10:05Z

Pushed: 2026-10-06T01:02:03Z

Default branch: main

Fork: no

Archived: no

README:

ALoDLM: Adaptively Looped Diffusion Language Models

Diffusion language models (DLMs) enable fast generation by predicting multiple tokens in parallel, yet their practical adoption remains hindered by a persistent quality gap relative to comparably sized autoregressive models. We attribute this gap to a *computation-difficulty mismatch*: within a partially observed sequence, some unknown tokens are readily predictable, while others require substantially more computation to resolve. Existing DLMs, however, apply uniform computational depth to every unknown position at each denoising step.

We introduce ALoDLM, which replaces this uniform computation with token-adaptive latent recurrence. At each denoising step, ALoDLM iteratively refines representations in latent space, allocating computation based on token difficulty. Tokens ready to commit are fed back as discrete context, while unresolved tokens retain and refine their latent states through additional recurrent passes. To learn both token prediction and computation allocation end-to-end, we formulate token-wise computation schedules as latent variables and derive a conditional negative evidence lower bound (NELBO).

We train ALoDLM at 1.7B and 8B parameter scales. Across eleven benchmarks, ALoDLM outperforms all evaluated DLMs and the corresponding autoregressive baselines in average benchmark score at both scales. ALoDLM combines superior generation quality with fast parallel decoding, establishing a strong quality-efficiency trade-off under optimized inference engines.

[Results](#results) · [Method](#method) · [Getting started](#getting-started) · [Training](docs/training.md) · [Evaluation](docs/evaluation.md) · [Limitations](#limitations)

Weights: ALoDLM-1.7B · ALoDLM-8B


*ALoDLM combines strong benchmark performance with fast parallel generation.*

Results

Across eleven benchmarks spanning reasoning, knowledge, mathematics, and code, ALoDLM achieves the highest average score among the evaluated diffusion models and exceeds the corresponding Qwen3 autoregressive baseline at both scales.

| Model scale | Qwen3 (AR) | Strongest compared DLM by average | ALoDLM | |---|---:|---:|---:| | 1.7B | 63.8 | 61.0 (SDAR) | 65.5 | | 8B | 78.5 | 75.1 (WeDLM) | 80.3 |

These are the paper's Table 1 results: the unweighted mean of eleven benchmark scores, expressed as percentages. Non-code tasks use accuracy or exact match; code tasks use execution-based pass@1. This quality comparison uses single-token commitment at each model pass.

• Fast parallel decoding. On GSM8K, ALoDLM-8B delivers approximately

2.7× the throughput of vLLM-served Qwen3-8B at comparable or higher accuracy, under single-stream inference on one NVIDIA B200.

• Quality per unit of computation. At the 93.25% accuracy threshold,

the paper reports 133.5 GFLOPs per generated token for ALoDLM versus 154.6 for WeDLM, a 13.6% reduction in estimated arithmetic cost.

• Test-time scaling and token adaptivity. Increasing the halting threshold

raises recurrent computation and improves average score in the reported sweep. Numerical tokens learn a lower first-pass halting probability, indicating a preference for further refinement without explicit token-difficulty labels.

The repository includes the ALoDLM optimized inference engine. See [optimized generation](#optimized-generation) for installation and usage.

Method


*Token commitments and latent refinement are interleaved inside each denoising step. A shared recurrent core reuses its parameters across passes. Entropy selects individual token commitments; the mean gate exit CDF over unresolved positions controls when the inner loop stops.*

Adaptive latent recurrence

A Qwen3 backbone is partitioned into a prelude, a shared recurrent core, and a coda. Each recurrent pass produces token predictions and learned halting probabilities:

1. Predict in parallel. Read out predictions for unresolved positions. 2. Commit confident tokens. Feed their token embeddings back as context. 3. Refine unresolved states. Continue from their accumulated latent representations until the halting criterion is met or the depth cap is reached.

Two complementary controls expose the quality-speed trade-off:

| Control | Role in parallel decoding | |---|---| | Entropy threshold tau | Determines which predictions can commit. A larger threshold permits less-confident commitments and generally favors faster generation. | | Cumulative halt threshold q | Determines when to end the inner loop. A larger threshold generally allows more latent refinement. |

The paper uses a maximum recurrent depth of four. At 8B, the recurrent core contains the middle 16 Transformer layers.

Learning where to spend computation

The paper treats token exit schedules as latent variables and derives a conditional negative evidence lower bound to jointly learn prediction and computation allocation. The implementation uses sequence-level outcome credit: each halting decision receives credit for its effect on the sequence as a whole.

Practical training combines a geometric depth prior, relaxed regularization of individual and average depth profiles, a first-pass control variate, weighted intermediate denoiser supervision, and an auxiliary next-token loss. The gate and backbone are trained jointly. See the [training objective](docs/method.md) for the complete implemented loss and its reductions.

Getting started

Clone this repository and choose one of the two inference environments below. The model weights are downloaded on first use; training is not required to try them. Use Python 3.12 for the commands below.

git clone https://github.com/amazon-science/ALoDLM.git
cd ALoDLM

Optimized generation

On Linux with an...