AFM Workshop, NeurIPS 2024Generation and refinement
COrAL: Order-Agnostic Language Modeling for Efficient Iterative Refinement
COrAL lets a language model predict and revise several token positions within a local window, making the trade-off between reasoning accuracy and decoding speed explicit.
- Result
- LogiQA accuracy 58.2% vs 55.1% for next-token greedy decoding, at 62.1 vs 33.6 accepted tokens/s (1.8×), both measured without KV cache at batch size 1.
- New vs prior
- Self-Refine refines through extra prompting passes, and Jacobi and look-ahead decoding parallelize forward decoding only; this work trains the model to predict forward and reconstruct backward within a sliding window.
- Limitation
- On HumanEval the full method lowers pass@1 from 64.6% to 13.0%, which the paper attributes mainly to syntax errors from inconsistent multi-token predictions.
Why should a token become permanent as soon as it is generated?
Autoregressive decoding commits to each token from left to right; when a later phrase makes an earlier choice look wrong, revising it takes another sequential pass. Parallel decoding cuts latency, but independently proposed tokens can disagree. COrAL keeps the sequential structure while making nearby predictions revisable.
The key restriction is locality: COrAL learns dependencies only within a bounded window, not arbitrary orderings over a whole sequence. A position can be predicted from the prefix or reconstructed with some later text visible, and agreement between these two views helps decide what to keep.
Closest prior work
Iterative refinement in LLMs usually runs on autoregressive models, either through prompting frameworks such as Self-Refine (Madaan et al., 2023) or by using feedback as a training signal, and each refinement adds sequential passes. Parallel decoding methods, such as speculative decoding with a draft model and Jacobi or look-ahead decoding, speed up generation but decode only in the forward direction, while non-autoregressive models often produce inconsistent outputs. COrAL trains the model itself to predict forward and reconstruct backward within a local window, so refinement happens inside the decoding loop.
Full COrAL gains on three of five benchmarks and fails on code
58.2 vs 55.1
LogiQA accuracy (%) · full COrAL vs next-token decoding
Throughput rises from 33.6 to 62.1 accepted tokens/s (1.8×). On GSM8K the same full setting gives 75.3 vs 74.1 at 1.1×.
13.0 vs 64.6
HumanEval pass@1 (%) · full COrAL vs next-token decoding
Turning off multi-forward prediction recovers 61.6, but at 28.8 vs 42.2 accepted tokens/s, slower than the baseline.
Five benchmarks, four decoding settings, one pair of axes
Each panel plots one benchmark from the paper: accuracy (pass@1 for HumanEval) against accepted tokens per second. Both axes start at zero and are the same in every panel. The arrow runs from the comparison setting to the highlighted one.
Next-tokenFull COrALNo verifierNo multi-forwardFilled accent: Full COrAL (highlighted). Dark: next-token (comparison).
- Accuracy (%)
- 74.1 → 75.3 (+1.2)
- Tokens/s
- 39.7 → 43.4 (1.09×)
- Accuracy (%)
- 21.8 → 22.7 (+0.9)
- Tokens/s
- 38.7 → 44.4 (1.15×)
- Accuracy (%)
- 55.1 → 58.2 (+3.1)
- Tokens/s
- 33.6 → 62.1 (1.85×)
- Accuracy (%)
- 63.2 → 62.7 (−0.5)
- Tokens/s
- 33.2 → 38.2 (1.15×)
- Pass@1 (%)
- 64.6 → 13.0 (−51.6)
- Tokens/s
- 42.2 → 45.8 (1.09×)
Horizontal: accepted tokens per second, 0–160. Vertical: accuracy or pass@1, 0–100%. Changes in parentheses are percentage points and throughput ratios, computed from the reported values.
Full COrAL vs Next-token: accuracy is higher on 3 and lower on 2 of five benchmarks. Throughput is 1.09× to 1.85× the Next-token rate. HumanEval pass@1 goes from 64.6 to 13.0 (−51.6 points).
All 40 reported values: quality in % and accepted tokens per second for each setting.
| Setting | GSM8K Table 1 | MATH Table 1 | LogiQA Table 2 | ReClor Table 2 | HumanEval §4.1 table | |||||
|---|---|---|---|---|---|---|---|---|---|---|
| Accuracy | Tok/s | Accuracy | Tok/s | Accuracy | Tok/s | Accuracy | Tok/s | Pass@1 | Tok/s | |
| Next-token | 74.1 | 39.7 | 21.8 | 38.7 | 55.1 | 33.6 | 63.2 | 33.2 | 64.6 | 42.2 |
| Full COrAL | 75.3 | 43.4 | 22.7 | 44.4 | 58.2 | 62.1 | 62.7 | 38.2 | 13.0 | 45.8 |
| No verifier | 72.4 | 156.8 | 20.0 | 139.7 | 55.7 | 99.1 | 61.6 | 72.0 | 6.5 | 119.0 |
| No multi-forward | 78.7 | 14.9 | 24.3 | 11.5 | 59.1 | 8.9 | 64.7 | 11.3 | 61.6 | 28.8 |
The paper attributes the HumanEval drop mainly to invalid syntax, 70.1% of samples in its no-verifier error analysis (Figure 5). The 61.6 without multi-forward prediction was obtained with a stricter acceptance threshold, ϵ = 0.5. Speeds were measured without a KV cache at batch size 1 on a single GPU, and code generation ran on different GPUs from the reasoning tasks, so compare speeds within a panel.
Measured: Table 1 (GSM8K, MATH), Table 2 (LogiQA, ReClor) and the §4.1 HumanEval table, Mistral-7B-v0.3, greedy next-token baseline, throughput in accepted tokens per second (hardware and ϵ in Appendix D). Derived: differences and ratios. Axes start at zero; points are reported experiments, not interpolations.
Among COrAL’s settings, backward refinement without multi-forward prediction gives the highest accuracy and the lowest throughput, and multi-forward prediction without verification the highest throughput and the lowest accuracy; the full method lies between them on every benchmark. The abstract’s 4.6-point GSM8K gain and 3.9× speedup thus come from two different settings. Throughput is accepted tokens per second, measured for both decoders without a KV cache at batch size 1.
Code is the boundary. The paper attributes the HumanEval drop mainly to invalid syntax; recovering most of the pass rate requires turning off multi-forward prediction and raising the acceptance threshold, which runs slower than the baseline. Local agreement helps refinement but does not ensure well-formed programs.
Scope. Order-agnostic training is applied only during task-specific fine-tuning of Mistral-7B-v0.3, with fixed windows of 4 forward and 8 backward tokens; extending it to pretraining is left to future work. Read the study ↗
Research summary
Iterative refinement has emerged as an effective paradigm for enhancing the capabilities of large language models (LLMs) on complex tasks. However, existing approaches typically implement iterative refinement at the application or prompting level, relying on autoregressive (AR) modeling. The sequential token generation in AR models can lead to high inference latency. To overcome these challenges, we propose Context-Wise Order-Agnostic Language Modeling (COrAL), which incorporates iterative refinement directly into the LLM architecture while maintaining computational efficiency. Our approach models multiple token dependencies within manageable context windows, enabling the model to perform iterative refinement internally during the generation process. Leveraging the order-agnostic nature of COrAL, we introduce sliding blockwise order-agnostic decoding, which performs multi-token forward prediction and backward reconstruction within context windows. This allows the model to iteratively refine its outputs in parallel in the sliding block, effectively capturing diverse dependencies without the high inference cost of sequential generation. Empirical evaluations on reasoning tasks demonstrate that COrAL improves performance and inference speed, respectively, achieving absolute accuracy gains of 4.6% on GSM8K and 4.0% on LogiQA, along with inference speedups of up to 3.9× over next-token baselines. Preliminary results on code generation indicate a drop in pass rates due to inconsistencies in order-agnostic outputs, highlighting the inherent quality–speed trade-off. Our code is publicly available at https://github.com/YuxiXie/COrAL.
Full paper ↗Citation BibTeX
@misc{xie2024coralorderagnosticlanguagemodeling,
title={COrAL: Order-Agnostic Language Modeling for Efficient Iterative Refinement},
author={Yuxi Xie and Anirudh Goyal and Xiaobao Wu and Xunjian Yin and Xiao Xu and Min-Yen Kan and Liangming Pan and William Yang Wang},
year={2024},
eprint={2410.09675},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2410.09675}
}