Label Collapse in Diffusion Language Model In-Context Learning

AI summary

Overview

Research area: Natural Language Processing — in-context learning (ICL), label bias, and discrete diffusion language models (dLLMs), with autoregressive (AR) models as baselines.

Technical level: Intermediate. The paper is readable with a working knowledge of prompting and logits, though the diagnostic machinery (log-odds separation, logit lens, calibration baselines) is mechanics-flavored.

Scope: A cross-architecture study of label collapse — a failure mode in which a model emits query-independent labels — evaluated on five 8–9B-scale models, 20 label word pairs, four demonstration arrangements, and demonstration counts from 4 to 30.

What This Paper Is About

Class-balanced ICL prompts are often built by simply alternating positive and negative demonstrations, on the assumption that balanced labels mean balanced behavior. The authors show this assumption fails: under an alternating arrangement, models can output the same label word for essentially every query regardless of input, and because such a model scores roughly 50% on a balanced binary task, standard accuracy reporting cannot tell it apart from a model that is genuinely balanced.

The paper's goal is to name this failure ("label collapse"), establish how widespread it is across both AR and diffusion architectures, and supply a diagnostic — the log-odds separation between the two label words at the answer position — that reveals collapse hidden beneath near-chance accuracy.

Key Contributions

1. Identification of label collapse as an architectural universal. The authors demonstrate collapse across five models — LLaDA-8B-Base, iLLaDA-8B-Base, LLaDA-8B-Instruct, Llama-3-8B-Instruct, and Qwen3.5-9B-Base — over 20 label word pairs and four arrangements, with severity ranging from total collapse to model-dependent graded failure.

2. A demonstration that accuracy is blind to the failure. A collapsed model producing one label for all queries lands near chance on a balanced binary task, which is numerically indistinguishable from balanced competence, so the failure can persist undetected in deployment under standard protocols.

3. Adoption of separation as a single-pass, model-agnostic diagnostic. Separation is the log-odds of the two label words at the query position (the same quantity label-bias calibration operates on), read from one forward pass without correction. The paper pairs it with a measurement protocol that audits for tokenizer splits (e.g., "toxic" as to + xic) and output-format drift, and shows the bias structure it exposes predicts which collapsed conditions calibration can and cannot repair.

4. Evidence that arrangement can reverse, but not avert, collapse. Arrangement flips which label a weak-bias model collapses to, while strong-bias collapse (circle/square, safe/unsafe) survives all four arrangements at every tested demonstration count.

Main Findings

• Collapse has a three-tier structure governed by the label words, not the architecture. Circle/square stays above 0.80 toward "circle" in 170 of 172 generation conditions with K ≥ 4, and the direction is fixed before any demonstration: at K = 0 the single-pass prior of all five models already leans toward "circle" (mean positive probability 0.70–0.83). Strong/weak is the only pair healthy across all five models.

• The two dLLM lineages have nearly disjoint dangerous-pair sets. At K = 12 under blocked, LLaDA-Base collapses on 6 of 10 antonym pairs (all toward the positive word: fast, kind, light, quiet, rich, wet), iLLaDA collapses on none (pos_rate 0.46–0.51), and LLaDA-Instruct reaches 7 of 10.

• A surprising inversion on the task's own label pair. On positive/negative, iLLaDA is the only healthy model: across 51 arrangement–K conditions with K ≥ 4 it never collapses (pos_rate 0.38–0.56, accuracy 0.87–0.96), while the other four lean toward "positive" (0.69–1.00 pooled, accuracy 0.48–0.52). The K = 0 priors do not explain this — iLLaDA's prior is 0.70, comparable to LLaDA-Base's 0.77 — so the repair is attributed to iLLaDA's pre-training corpus (12T tokens vs. LLaDA's 2.3T), not to chat SFT.

• Hard collapse ignores arrangement entirely. Safe/unsafe stays pinned to "unsafe" for iLLaDA (pos_rate 0.000 in all 51 conditions) and Qwen3.5 (≤ 0.03 in all 24) across all four arrangements and every tested K. Adding demonstrations does not merely fail to help — iLLaDA's K = 0 prior leans "safe" (0.61), and demonstrations flip it to "unsafe."

• Both AR models collapse too, with distinct signatures. At K = 8 alternating, Llama-3 collapses on 13 of its 18 parseable pairs and Qwen3.5 on 7 of 20. Qwen3.5 escapes the label vocabulary on right/left (out-of-pair rate 45–74% at small K) before returning to balance within it. Llama-3's true/false separation oscillates by parity (even K: +3.0 to +4.0; odd: near zero) and generation follows (even K ≥ 10: 0.79–0.91 toward "true"; odd: 0.28–0.44).

• The bias is set at or before the final projection layer. Applying the logit lens to all 33 layers (L0–L32) of LLaDA-Base, iLLaDA, and Llama-3 shows separation near zero through intermediate layers and emerging at L29–L32 for Llama-3 and iLLaDA. LLaDA-Base is heterogeneous: clean/toxic, safe/unsafe, true/false, and circle/square already show large separation at layer 0, consistent with token-embedding-level bias.

• Roughly three quarters of collapse directions trace to a token-level prior. Zero-shot separation sign agrees with high-K collapse direction in 33 of 45 collapsed model–pair combinations (73%); empty-context unigram logits agree in 28 of 37 (76%). A pure word-frequency account matches in only 7 of 14 cases where the two label words differ by at least a factor of 2.5 in corpus frequency.

• Separation predicts behavior with high precision. Across 1,502 paired conditions, the criterion |sep| > 3 identifies behavioral collapse with 92.5% precision and 84.7% recall. Generation and single-pass pos_rate agree within 0.10 in 88% of conditions and within 0.30 in 95%.

• One systematic decoupling, and it is SFT-dependent. LLaDA-Instruct on true/false collapses in generation (0.69–0.99 toward "true") while its single pass retains query dependence (class accuracy 0.65–0.94, mean separation only −1.9 to +0.03) — read as iterative denoising amplifying a direction the single pass slightly opposes. LLaDA-Base shows only a transient version of this (generation 0.04–0.20 at K = 5–7 while the single pass stays 0.89–0.91 accurate), vanishing from K = 8.

• Arrangement can reverse collapse direction. On ancient/modern every model flips: alternating collapses to "ancient" (0.93–1.00 for three models), grouped collapses all five to "modern" (0.07–0.17). On cat/dog, LLaDA-Instruct collapses to "cat" under alternating (0.82–1.00) and to "dog" under grouped (0.00–0.17). No arrangement is uniformly safest: grouped is more severe on average than alternating for LLaDA-Base (0.43 vs. 0.35), iLLaDA (0.35 vs. 0.25), and Qwen3.5 (0.33 vs. 0.28), while alternating is slightly more severe for Llama-3 (0.36 vs. 0.32) and LLaDA-Instruct (0.37 vs. 0.33).

• Calibration results locate the bias rather than fix it. At K = 8 on the single-pass channel (86 conditions, 45 collapsed), contextual calibration repairs 16 of 45 and query-conditional PMI repairs 4 of 45, while batch calibration — which subtracts the test-batch mean — returns all 45 to balance and restores accuracy in 36. That 45/45 is described as the formal signature of query-independent bias; batch calibration is transductive, so it diagnoses rather than deploys.

Methodology in Plain English

The authors build a controlled comparison where the only things that change are the label words and how the same demonstrations are ordered.

Prompts. Every prompt lists K labeled examples ("sentence … Answer: <label>") and ends with a query and a bare "Answer:" slot. Each condition uses 200 query sentences. Demonstrations come from SST-2 (sentiment) and XSTest (safety), and the 20 label word pairs span sentiment, safety, geometric, and abstract/antonym domains — ten "task pairs" like positive/negative and safe/unsafe, and ten "antonym pairs" like hot/cold and rich/poor that reuse the same demonstration pool with only the label words swapped. That swap makes label-word identity the sole manipulated variable.

Arrangements. Four orderings of the same class-balanced multiset (⌈K/2⌉ positives, ⌊K/2⌋ negatives): alternating (PNPN), grouped (PPNN), blocked (up to three per block, e.g., PPPNNNPN at K = 8), and shuffled (a fixed-seed random permutation). Demonstration counts run from 4 to 30, with different exact grids per model for compute reasons; all cross-model snapshots use K = 8, the one value evaluated for every model.

Two measurement channels. The generation channel reads what the model actually outputs and records pos_rate (fraction matching the positive label) and accuracy. The single-pass channel takes one forward pass and reads the logits of the two label words at the answer position; their difference is the separation, which equals the log-odds of the two label words because the softmax denominators cancel.

Severity grades. Four explicit tiers defined on pos_rate and accuracy: hard collapse (pos_rate ≤ 0.05 or ≥ 0.95 with accuracy ≤ 0.55, in every tested condition), collapse (≤ 0.20 or ≥ 0.80 with accuracy ≤ 0.60), drift (beyond ±0.15 of balance with accuracy ≥ 0.70), and healthy (within ±0.15 with accuracy ≥ 0.85). The thresholds are anchored to the observed distribution rather than tuned — the 0.05 boundary sits just above 0.03, the largest pos_rate found in any hard-collapse condition — and shifting any boundary by ±0.05 reclassifies at most 7.0% of conditions.

Localizing the bias. They run the logit lens across all 33 layers to see where separation appears, then measure the label prior with no demonstrations at all in two independent ways (query + answer slot, and a bare answer slot), and cross-validate the two. They also test whether existing calibration methods can undo the bias, and check whether separation predicts generation behavior.

Measurement hygiene. Because tokenizers may split a label word into several tokens and outputs may drift off the answer line, every condition is audited for both artifacts, and multi-token cases are handled by parsing at the text level rather than by first-token matching.

Why This Matters

The paper's central practical claim is that a model can look "balanced" on a balanced benchmark while ignoring its input entirely, and that no accuracy-based protocol will notice. Because the diagnostic is a single forward pass with no retraining and no reference-point correction, it is cheap enough to run as a routine audit. The finding that collapse direction is largely fixed by the label words and by training lineage — and is arrangement-invariant for strong-bias pairs — means the risk is predictable from the label set and a K = 0 probe rather than discovered in production.

Real-world applications:

• Benchmark auditing and model cards: Flagging evaluations where a reported ~50% binary accuracy may reflect a model that always answers one way.
• Prompt and label-schema design: Choosing label words (e.g., avoiding pairs like circle/square) that do not trigger query-independent output, and knowing that arrangement alone cannot rescue a strong-bias pair.
• Safety and content classification: Safe/unsafe and clean/toxic are exactly the pairs where hard collapse appears, so a deployed safety classifier could be returning a constant label while appearing calibrated.
• Calibration in deployment: Distinguishing biases that a reference-point correction can remove from those it cannot, and recognizing that batch calibration's success is a diagnostic signature rather than a deployable fix.

Industry relevance: Any team shipping few-shot classifiers or evaluating dLLMs should care. The paper covers diffusion models specifically because their ICL is new and their failure modes are uncharacterized; the safe/unsafe result — iLLaDA pinned at pos_rate 0.000 across all 51 conditions — is a concrete warning for safety-critical deployments. The lineage findings also matter for model selection: chat SFT introduced new failure modes (a stable third-word "neutral" answer, direction reversal, generation–single-pass decoupling) without repairing the positive/negative collapse, while a different pre-training procedure repaired it entirely.

Future Directions

• How iterative denoising interacts with label biases. The authors state this remains open. The LLaDA-Instruct true/false case — a discriminative single pass paired with a collapsed generation — suggests denoising can amplify a direction the forward pass opposes, but the mechanism is not resolved.

• Probing beyond the output head. The paper applies the logit lens but notes that extending trained linear probes and denoising-step trajectories to dLLMs remains open, which would show what the intermediate layers encode during denoising rather than only what the output projection emits.

• Turning diagnosis into repair. Batch calibration returns all 45 collapsed single-pass conditions to balance but is transductive. Injecting the correction into dLLM decoding repaired isolated conditions (LLaDA-Base positive/negative: 0.89 → 0.51) but left pinned pairs like circle/square and true/false unmoved, with occasional overshoot, so a non-transductive fix is unresolved.

• Generalizing beyond binary tasks. The conclusion points to a preregistered four-class extension reported in Appendix J, where the severity spectrum remains word-set-governed and batch calibration retains near-universal repair while contextual calibration fails.

Target Audience

Researchers working on in-context learning, prompt sensitivity, and evaluation methodology; dLLM developers who need to know whether AR label-bias findings transfer to iterative denoising; and practitioners deploying few-shot classifiers or safety filters who need a cheap pre-deployment check for query-independent output. Readers focused on scaling laws or architectural novelty will find less here — the paper deliberately holds scale (8–9B) and task type (binary classification) fixed to isolate label-word and arrangement effects, and it notes that which pairs are dangerous is a property of the training procedure within the models and scale tested.

Authors’ 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.