Label Collapse in Diffusion Language Model In-Context Learning

Ziye Guo; Ye Yuan

cs.CLOct 12, 2026

Label Collapse in Diffusion Language Model In-Context Learning

Authors: Ziye Guo, Ye Yuan

Organizations: University of Connecticut · 2McGill University

Abstract

In-context learning (ICL) applies large language models to new tasks by prepending a few labeled demonstrations to a query. Prior work has studied how demonstration ordering affects ICL, but the interaction between label arrangement patterns and model biases has not been systematically examined. We define label collapse as the failure mode in which the model generates query-independent labels, and show that it is architecturally universal at the tested 8-9B scale: five models spanning autoregressive (Llama-3, Qwen3.5) and discrete diffusion (LLaDA, iLLaDA, LLaDA-Instruct) families all exhibit label collapse even under the class-balanced alternating arrangement, with some label word pairs collapsing in virtually every tested condition and others flipping their collapse direction with arrangement. Critically, accuracy alone cannot detect this collapse: models can achieve near-chance accuracy while generating query-independent labels, a score indistinguishable from balanced performance, so the failure can persist undetected in deployment. We measure separation, the log-odds of the two label words at the query position (the quantity that label-bias calibration operates on), as a complementary diagnostic that exposes collapse hidden beneath surface accuracy, together with a measurement protocol that guards label extraction against tokenizer and formatting artifacts. For autoregressive models, collapse persists where the majority-label-bias account does not apply, and arrangement patterns can reverse the collapse direction; for diffusion models, we present a systematic evaluation of ICL failure modes, complementing contemporaneous studies of dLLM context comprehension and demonstration-order sensitivity, and identify open questions about how iterative denoising interacts with label biases.

Figures & tables

Figure 1: K-sweep generation pos_rate for the two anchor pairs (circle/square, safe/unsafe; alternating arrangement). Gray dashed line = 0.5 (balanced). circle/square collapses to 1.0 for all models at every $K$ ; on safe/unsafe, iLLaDA collapses to 0.0 while Llama-3 gradually recovers toward 0.5 as $K$ grows and LLaDA-Base reverses direction at high $K$ . The remaining two representative pairs (positive/negative, clean/toxic) are shown in Appendix B .

Universal-collapse
Word pair	LLaDA	iLLaDA	LLaDA-Inst.	Llama-3	Qwen3.5
circle/square	1.00 ✗	1.00 ✗	1.00 ✗	1.00 ✗	1.00 ✗
Healthy (all five models functional)
strong/weak	0.49	0.48	0.47	0.59	0.48
Model-dependent
positive/negative	0.89 ✗	0.50	0.90 ✗	0.91 ✗	0.88 ✗

Table 1: Generation pos_rate for all 20 word pairs at

K{=}8

under alternating (200 queries per cell), sorted into severity tiers. Tier assignment uses all arrangements and

K

values tested ( not only this snapshot): “universal” collapses under every arrangement/

K

; “healthy” stays functional for all five models. ✗ = collapse (

|\text{{pos\_rate}}-0.5|\geq 0.3

);

\triangle

= drift (

0.15\leq|\text{{pos\_rate}}-0.5|<0.3

); † = more than 30% of outputs leave the label vocabulary or the answer format; such cells are excluded from collapse counts (Appendix D ). The ancient/modern row reflects the alternating arrangement only: under grouped and shuffled its cells reverse toward “modern”, and under blocked the two dLLM base models return to balance (Section 6 ).

Figure 2: Generation vs. single-pass (sp) pos_rate on the same 200 queries per condition (alternating arrangement). (a) LLaDA-Base, safe/unsafe, $K{=}4$ – $30$ : the two channels track each other point by point (max difference $0.13$ ). (b) LLaDA-Instruct, true/false, $K{=}4$ – $30$ : generation collapses to “true” ( $0.69$ – $0.99$ ) while sp does not ( $0.15$ – $0.48$ ); sp nonetheless retains classification accuracy of $0.65$ – $0.94$ , so the single pass keeps query dependence that iterative denoising erases.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Model	PNPN	PPNN	PPPNNN	shuffled
LLaDA-8B-Base	0–30 (all)	4–14 (all), 16, 30	6, 8, 9, 12, 15, 16, 18, 30	8, 16, 30
iLLaDA-8B-Base	0–30 (all)	4–14 (all), 16, 30	6, 8, 9, 12, 15, 16, 18, 30	8, 16, 30
LLaDA-8B-Instruct	0–13, 15, 16, 20, 24, 30	4, 8, 12, 16, 30	8, 12, 30	8
Llama-3-8B-Instruct	0–30 (all)	4, 8, 12, 16, 30	8, 12, 16, 30	8, 16, 30
Qwen3.5-9B	0–13, 15, 16, 20, 24, 30	4, 8, 12, 16, 30	8, 12, 30	8

Appendix

Table 2: Complete demonstration-count grids per model and arrangement: all

K

values actually evaluated, unioned over word pairs. The cross-arrangement comparisons in the text use the common balanced subset

K\in\{4,8,12,16,30\}

Figure 3: K-sweep generation pos_rate for the remaining two representative pairs (positive/negative, clean/toxic; alternating arrangement; gray dashed line = 0.5). Complements Figure 1 : on positive/negative the other four models lean toward “positive” while iLLaDA stays balanced; on clean/toxic LLaDA-Base and LLaDA-Instruct collapse toward “clean” while iLLaDA and Qwen3.5 stay near balance.

Figure 4: safe/unsafe under alternating, per model, as a function of $K$ : generation pos_rate (left) and single-pass separation (right). The two panels track each other closely (the parity oscillation of LLaDA-Base is visible in both), supporting the use of separation as a predictor of generation in this regime. The shaded band marks the $\lvert\mathrm{sep}\rvert\leq 3$ criterion. LLaDA-Instruct and Qwen3.5 were evaluated at the sparser $K$ grid of Section 2 .

Figure 5: Arrangement comparison at $K{=}8$ (generation pos_rate ; four representative pairs, five models). circle/square collapses under all four arrangements for all models; safe/unsafe collapses for iLLaDA and Qwen3.5 under all arrangements, while Llama-3 and LLaDA-Instruct show arrangement-dependent severity and LLaDA-Instruct cat/dog reverses direction between alternating and grouped. Complements Table 3 with the four representative pairs of Figures 1 and 3 .

Model	Word pair	PNPN	PPNN	PPPNNN	Pattern
LLaDA-Base	safe/unsafe	0.35 $\triangle$	0.02 ✗	0.24 $\triangle$	Grouped hardest
iLLaDA-Base	safe/unsafe	0.00 ✗	0.00 ✗	0.00 ✗	All collapse
LLaDA-Inst.	true/false	0.96 ✗	0.72 $\triangle$	0.83 ✗	Always “true”-biased
LLaDA-Inst.	cat/dog	0.82 ✗	0.17 ✗	0.47	Direction reversal
Llama-3	cat/dog	1.00 ✗	0.89 ✗	0.58 †	Alternating collapse
Qwen3.5-9B	safe/unsafe	0.00 ✗	0.00 ✗	0.00 ✗	All collapse

Appendix

Table 3: Cross-arrangement generation pos_rate at

K{=}8

(SST-2 demonstrations for cat/dog, XSTest for safe/unsafe and true/false); column heads follow the arrangement codes of Section 2 . “Collapsed” cells (

|\text{{pos\_rate}}-0.5|\geq 0.3

) are marked ✗. † Llama-3 cat/dog under blocked was not run at

K{=}8

; the value shown is at

K{=}12

(

0.57

;

K{=}30

gives

0.59

\triangle

marks drift (

0.15\leq|\text{{pos\_rate}}-0.5|<0.3

Figure 6: Distribution of per-condition gen–sp gaps. Empirical CDF of the $|$ gen $-$ sp $|$ pos_rate gap over all 1,502 paired conditions (pairing of Appendix B ; generation readings use the re-parsed values of Appendix D ); thin curves per model, thick black pooled. 88% of conditions lie below $0.10$ and 95% below $0.30$ (dashed guides, pooled fractions marked). The tail concentrates in LLaDA-Instruct: 36 of its 256 conditions lie at or above $0.30$ , versus 3–16 for each other model, and its curve is the lowest at the $0.30$ guide. Figure 7 traces one cell of each kind.

Figure 7: $K$ -trajectories of one alternating-arrangement cell per kind of above-threshold gen–sp behavior (blue solid: generation; red dashed: single pass; gray dashed line: balance; gray band in (d): outputs matching neither label word). (a) Channel flip (Llama-3 right/left): at even $K$ the two channels point to opposite labels (gap up to $0.91$ ) while the single pass sits at chance accuracy ( $0.50$ ), having itself lost query information. (b) Transient episode (LLaDA-Base true/false): generation dips to $0.04$ – $0.20$ at $K{=}5$ – $7$ while the single pass stays accurate ( $0.89$ – $0.91$ ), rejoining it from $K{=}8$ onward. (c) Persistent exception (LLaDA-Instruct true/false): generation at $0.69$ – $0.99$ at every $K$ while the single pass ( $0.65$ – $0.94$ accurate) stays mid-range. (d) Invalid outputs (iLLaDA right/left): the gray band marks outputs matching neither label word (37–80% of queries); among the parseable remainder, generation accuracy is $0.90$ – $1.00$ , so its low pos_rate reflects off-label emission rather than a query-blind label choice.

Figure 8: Single-pass pos_rate before (solid) and after (dashed) contextual calibration (alternating; panels follow the word-pair convention of Figure 1 ). Left, positive/negative: all four collapsed models move to balance ( $0.49$ – $0.62$ ; iLLaDA is balanced already, and LLaDA’s repair decays beyond $K{=}20$ ). Right, safe/unsafe: Llama-3 is lifted steadily toward balance and Qwen3.5 at large $K$ , while iLLaDA stays pinned and LLaDA overshoots at small $K$ .

Figure 9: Word pairs that calibration barely moves (single-pass pos_rate , alternating; solid before, dashed after). Left, circle/square: LLaDA, iLLaDA, and Qwen3.5 stay at $1.0$ at every $K$ ; Llama-3, the boundary case, recovers only at part of the grid ( $0.49$ at $K{=}8$ ). Right, true/false: no collapsed model is repaired (iLLaDA stays at $0.0$ , Llama-3 stays above $0.9$ ); the other three models sit near balance already.

Figure 10: Generation pos_rate before (solid) and after (dashed) injecting the contextual-calibration correction into dLLM decoding (alternating; three dLLMs, word-pair panels as in Figure 1 ). Left, safe/unsafe: the correction shifts LLaDA-Base partway toward balance (overshooting at $K{=}4$ ) and lifts iLLaDA transiently off the floor ( $0.255$ peak at $K{=}12$ ). Right, true/false: decoding is essentially unchanged. The injection repairs margins, not the hard tier.

Figure 11: Generation pos_rate on the three equal-length pairs of Appendix F (alternating arrangement, 200 queries, $K{=}4$ – $30$ ). Left: new/old collapses toward “old” within $K{\leq}8$ for LLaDA-Base, LLaDA-Instruct, Llama-3, and Qwen3.5, while iLLaDA alternates with parity throughout (even $K$ at “new”, odd $K$ at “old”). Middle: moral/immoral collapses toward “immoral” for four models while iLLaDA stays functional. Right: valid/invalid shows a parity-structured drift in LLaDA-Base, iLLaDA, and Qwen3.5 (odd $K$ toward “invalid” from the lowest $K$ , even $K$ rising toward “valid” as $K$ grows, Qwen3.5 joining only at $K{\approx}20$ ), a mild “invalid” lean for LLaDA-Instruct, and near-balance for Llama-3. The gray dashed line marks $\texttt{pos\_rate}=0.5$ .

Model	A1	A2	A3	W1	U1	U2	U3	U4	U5	U6	V1	V2
LLaDA	0/15	3/15	8/15	7/15	15/15	10/15	10/15	15/15	10/15	10/15	11/15	9/15
LLaDA-Inst.	8/15	9/15	9/15	5/15	10/15	11/15	10/15	15/15	6/15	10/15	14/15	11/15
iLLaDA	1/15	0/15	5/15	0/15	13/15	10/15	10/15	15/15	10/15	9/15	11/15	10/15
Dream-Base	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*
Dream-Inst.	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*	0/15*
Llama-3	0/15	—	—	—	10/15	—	—	—	—	—	—	6/15

Appendix

Table 4: Multi-class collapse spectrum: collapsed generation conditions (threshold

0.70

) over

K\geq 12

, main demonstration pool, by label set. Dream rows: all cells are out-of-set collapse (out-of-set rate

1.000

), not health (text below). AR rows cover A1/U1/V2 only.

In-Context Learning

In-context learning studies how models adapt to new tasks from examples or instructions supplied at inference time, without changing their weights.

See all 386 papers

From Copying to Composition: How Language Models Develop In-Context Learning
Oct 13, 2026Minglai Yang, Xinyu Guo, Renxiong Wang +3In-Context Learning Compositional Generalization
Scale AI · University of Arizona
In-Context Collapse in Vision-Language Models and How to Mitigate it?
Aug 3, 2026Mohammad RostamiVLM Robustness Vision-Language Models
Amazon Generative AI Innovation Center∗, USA.
Where to Place the Query? Unveiling and Mitigating Positional Bias in In-Context Learning for Diffusion LLMs via Decoding Dynamics
Apr 26, 2026Zhengheng Li, Panrui Li, Xuyang Liu +1In-Context Learning Diffusion Language Model Decoding
Southeast University

Oct 13, 2026J/K move · Enter open · S save

From Copying to Composition: How Language Models Develop In-Context Learning

Minglai Yang, Xinyu Guo, Renxiong Wang, Sunjie Hou+2

Scale AI · University of Arizona

In-Context Learning Compositional Generalization

Open paper AskPDF arXiv

In-context learning (ICL) allows language models to adapt their behavior from instructions and demonstrations in a prompt, but accuracy alone does not reveal how contextual information controls their predictions. We study the acquisition and trainability of ICL with a controlled diagnostic that separates copying a nearby demonstration from integrating information across the prompt. Across 17 models, about a third prefer the copied answer, a few are undecided, and the rest prefer the rule-consistent answer, which goes with higher accuracy and with control by the requested value rather than the closest conflicting example. Checkpoint analyses show that models begin prioritizing task-relevant prompt tokens over conflicting examples long before their accuracy improves beyond chance level. In matched training experiments, synthetic documents that define endings and then query them shift models toward compositional responses on the controlled diagnostic, whereas fine-tuning on definition exercises built from rewritten real documents improves accuracy when prompts reassign label meanings. The results suggest that ICL is not acquired as a single, uniform capability. Instead, prompt control, answer preference, construction without candidates, and transfer can develop along distinct trajectories.

Abstract

In-Context Learning

Explore similar work

From Copying to Composition: How Language Models Develop In-Context Learning

In-Context Collapse in Vision-Language Models and How to Mitigate it?

Where to Place the Query? Unveiling and Mitigating Positional Bias in In-Context Learning for Diffusion LLMs via Decoding Dynamics