From Confusion to Clarity: Confusion-Aware Retrieval and Knowledge Injection for Text Classification

AI summary

Overview

Research area: Natural Language Processing — zero-shot and retrieval-augmented LLM text classification over large, semantically overlapping label taxonomies. Technical level: Intermediate (comfortable with retrieval-augmented prompting, Macro F1, confusion matrices, and cross-model transfer concepts). Scope: The paper proposes a three-phase, fine-tuning-free framework — confusion-aware candidate retrieval, offline generation of pairwise disambiguation rules, and single-call rule injection at inference — and evaluates it on three classification benchmarks and across seven models from 2B to 32B parameters.

What This Paper Is About

Large language models struggle to classify text when a taxonomy contains hundreds of labels that overlap in meaning, because the distinctions are domain-specific and are not captured during pre-training. The standard fix — retrieve the top-K labels by embedding similarity and let the model choose among them — shrinks the candidate set but gives the model no signal for telling two similar candidates apart. The goal of this work is to identify the exact label pairs the model confuses, make sure both members of those pairs appear as candidates, and supply targeted rules that explain how to choose between them, all without fine-tuning.

Key Contributions

1. A candidate augmentation strategy that applies confusion-matrix analysis to the baseline classifier's training errors, so that systematically confused labels co-occur in the retrieval set as "confusion partners."
2. A three-stage knowledge generation pipeline — per-example reasoning, consolidation, and bidirectional merge — that produces one pairwise disambiguation rule per confused label pair, forming a knowledge base.
3. Empirical validation on three datasets and seven models spanning 2B–32B parameters, showing consistent Macro F1 gains and demonstrating that a knowledge base generated once by a large model transfers to smaller, cheaper classifiers without retraining.
4. A robustness comparison against fine-tuned ModernBERT-base and RoBERTa-base on Flipkart L3 under clean data, test-time perturbations, label noise, and data scarcity, plus a human evaluation of generated rule quality.

Main Findings

• Best Macro F1 in every column of the main comparison. With Qwen3-32B, the full approach reaches 58.9 (LEDGAR), 81.5 (WOS L1), 57.7 (WOS L2), 92.5 (Flipkart L1), 84.4 (Flipkart L2), and 76.7 (Flipkart L3), against retrieval at K=20 (56.3, 78.4, 52.4, 91.0, 81.8, 72.3) and zero-shot with all labels (57.1, 80.5, 56.6, 84.4, 73.3, 65.5).
• Largest gain on the deepest, most confusable level. Flipkart L3 improves by +11.2pp over zero-shot and +4.4pp over the strongest baseline (retrieval at K=20). The abstract reports gains of up to 10.0pp over retrieval baselines.
• Statistical significance is partial. Paired bootstrap (B = 10,000) shows significant gains on LEDGAR (+1.2, p = 0.016), WOS L1 (+1.2, p < 0.001), and Flipkart L3 (+2.7, p = 0.005). On WOS L2 (+0.3, p = 0.600), Flipkart L1 (+0.8, p = 0.194), and Flipkart L2 (+0.6, p = 0.360), the method remains best but the margin is within sampling noise.
• The two mechanisms are complementary. On Flipkart L3 at K=10, retrieval misses the correct label 14% of the time; confusion partners raise recall from 86% to 98.1% but push LLM miss errors from 15.5% to 21.9%; knowledge injection then pulls LLM misses back down to 17.1%.
• Confusion partners beat retrieving more labels. On Flipkart L3, partners at K=10 (18 candidates on average) yield 74.2 F1 and 98.1% recall, whereas plain retrieval needs K=20 (20 candidates) for only 72.2 F1 and 92.4% recall, and K=50 to match the recall.
• Pairwise rules beat per-label descriptions. On Flipkart with retrieval plus confusion partners, pairwise knowledge gives 92.5/84.4/76.7 F1 at L1/L2/L3, versus 91.4/83.5/74.8 for per-label knowledge and 90.8/81.9/74.2 with no knowledge; LLM miss rates follow the same ordering (3.8/10.0/17.1 versus 4.8/11.3/18.5 versus 6.2/13.3/21.9).
• Cross-model transfer works. With all knowledge generated by Qwen3-235B, small models (2–4B) gain +11.5pp on Flipkart L3 over retrieval-only and medium models (8–20B) gain +11.1pp; LEDGAR gains are +6.3pp and +5.3pp respectively.
• A fixed knowledge base also transfers. Reusing one Qwen3-235B knowledge base unchanged improves Macro F1 by 6.6 to 10.3 points over retrieval alone across Ministral-3B, Ministral-8B, GPT-OSS-20B, and Qwen3-32B.
• Fine-tuning wins on clean data but loses under production stresses. On Flipkart L3, ModernBERT and RoBERTa score 80.0 and 81.7 under clean data versus 76.7 for this framework; but with 20% label noise the framework holds 72.1 versus 60.0 and 61.0, with 10% of training data it reaches 70.2 versus 23.2 and 31.8, and across nine test-time perturbations it loses 5.0 points on average versus 17.0 and 9.0.
• Generated rules are mostly judged correct. Three authors each reviewed 50 sampled Flipkart rules from Qwen3-32B, producing 150 judgments, and accepted 87.3% of rules on average.

Methodology in Plain English

The framework runs in two phases. In the offline phase, the researchers first run a baseline retrieval classifier over the training split and record its mistakes. Every time the true label is predicted as another label, they increment a confusion count for that directed pair, rank pairs by frequency, and select the smallest set of pairs whose cumulative count covers at least a chosen percentage of all training errors (75% in the main experiments). For each label they also define a "confusion partner": the label with the highest combined count of being mistaken for it and vice versa. For each selected pair, a three-stage LLM pipeline produces a rule — the model first explains, per misclassified example, what signal indicates the correct label rather than the wrong one; it then consolidates those noisy observations into one directional rule; and if both directions of a pair exist, it merges them into a single symmetric rule.

At inference, the system retrieves the top-K candidates by embedding similarity (K=10), adds the confusion partner of every retrieved label, selects the rules from the knowledge base that apply to the resulting candidate set, injects them into the prompt alongside the labels, and makes a single classification call. All LLM calls use temperature 0. The main experiments use Qwen3-32B as both knowledge generator and classifier, with Qwen3-Embedding-8B for retrieval; the transfer experiments use Qwen3-235B as the generator and Ministral 3B/8B, Qwen3.5-2B/4B/9B, and GPT-OSS-20B as classifiers. WOS and Flipkart are hierarchical (2 and 3 levels) while LEDGAR is flat.

Why This Matters

Impact on research: The paper reframes label confusion as an error-analysis problem rather than a retrieval-scale problem, showing that the specific pairs a model gets wrong can be mined from its own confusion matrix and corrected with generated pairwise rules. It provides a deployment-oriented counterpoint to fine-tuning, showing where parameter-free inference-time context holds up better (label noise, scarce data, distribution shift) and where it does not (clean, abundant, stationary data).

Real-world applications:

• Product catalog classification, where catalogs have hundreds of near-synonymous categories such as "Women's Western Wear" versus "Women's Ethnic Wear."
• Customer service ticket routing, where misclassification degrades queue assignment.
• Legal document taxonomies and compliance filtering, evaluated here through the LEDGAR legal benchmark.
• Search relevance, which the authors identify as a downstream casualty of misclassification in large taxonomies.

Industry relevance: The method requires no fine-tuning, targets teams that access foundational models through managed APIs rather than running their own training and serving infrastructure, and lets a large model produce knowledge offline that a small, low-latency model consumes at inference. The measured cost is explicit: offline generation consumes 3.26M to 15.52M tokens once per taxonomy, and online median latency rises from 379–429 ms for retrieval-only to 879–1,238 ms for the full approach, with input tokens rising from roughly 1.2K–1.4K to 7.1K–10.1K per query.

Future Directions

• Testing whether the balance between pair coverage and prompt length holds for multilingual inputs, substantially larger taxonomies, and temporal drift, which the authors state they have not examined.
• Building automatic verification for generated rules, since the Consolidation and Bidirectional Merge stages reduce noise but do not guarantee correctness, and the 87.3% human acceptance rate still leaves room for error in high-stakes settings.
• Deciding when to re-run the offline phase, since rules can only cover confusion pairs observed in baseline errors and a changed taxonomy or changed confusion pattern requires regenerating affected pairs.
• Choosing between fixed and model-specific knowledge bases: fixed reuse transfers across classifiers, but discovering pairs per classifier improved results for three of four tested classifiers, so the trade-off between one-time reuse and per-model error analysis remains open.

Target Audience

Applied ML and NLP engineers building classification systems over large, fine-grained taxonomies who need inference-time alternatives to fine-tuning; researchers working on retrieval-augmented generation, prompt optimization, and knowledge injection for label disambiguation; and product or platform teams operating catalog, ticketing, or compliance classifiers where robustness to noisy labels and shifted inputs matters more than peak accuracy on clean data.

Authors’ abstract

Large language models (LLMs) struggle to classify text into taxonomies with many semantically similar labels, as the distinctions are domain-specific and not captured by pre-training. To handle large label spaces, a common approach retrieves top-$K$ candidate labels by embedding similarity and prompt the LLM to choose among them. However, top-$K$ retrieval reduces the number of candidates but does not help the model tell similar ones apart. When two similar labels both appear as candidates, the model lacks the signal to choose correctly between them. We propose a framework that (1) identifies which label pairs the model struggles to distinguish, (2) expands the candidate set to include confusable labels, and (3) generates targeted rules to differentiate between similar candidates. The framework requires no fine-tuning, and the generated rules transfer to smaller, cheaper models. On three benchmarks (WOS, Flipkart, LEDGAR), our approach improves Macro F1 by up to 10.0pp over retrieval baselines, with smaller models (2B--20B) gaining up to 11.5pp via cross-model transfer.