Research map / Label-efficient and robust learning / Recognition and long-tail

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification

ICML 2024 · Label-efficient and robust learning › Recognition and long-tail · approach: Multilingual language representation pretraining

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes OAK, an extreme classification framework that enriches document embeddings with trainable auxiliary knowledge piece embeddings linked by a linker module, trained with triplet and calibration losses.

Problem

Extreme classification on short-text documents is hard because sparse text representations are ambiguous, and existing XC methods ignore auxiliary data that is often available alongside documents.

Approach

OAK stores learnable embeddings for auxiliary knowledge pieces in a knowledge bank, uses a linker (an XC model) to retrieve relevant pieces at inference, fuses them with the document embedding via cross-attention, and adds a mutual-information calibration loss inspired by DPO so enriched representations beat plain ones.

Key results

• OAK reaches 33.71 P@1 on LF-WikiSeeAlsoTitles-320K vs 32.64 for NGAME, and 85.23 vs 84.01 on LF-Wikipedia-500K.
• On LF-ORCAS-800K, OAK gets 75.25 P@1 vs 72.47 for ANCE (about 3 points).
• In sponsored search offline evaluation, OAK reaches 68.80 Recall@100 vs 63.16 for the best production model M3, and online A/B testing gave a 0.84% CTR increase.

Measured effects reported in the paper

Surprising findings

• Using oracle ground-truth AKPs instead of linker-predicted AKPs gave lower accuracy on full-text datasets, attributed to missing links in the ground truth that the linker recovers.

Limitations

• Relies on a linker trained on document-AKP links, so noisy or missing AKP links can hurt results.
• Gains over NGAME on public benchmarks are modest (about 1-2 points on most datasets).
• Adds training cost (about 1.5x) and inference latency, and the AK bank requires a fixed pool of auxiliary items.

Experimental setup

Datasets: LF-WikiSeeAlsoTitles-320K LF-WikiSeeAlso-320K LF-WikiTitles-500K LF-Wikipedia-500K LF-ORCAS-800K SponsoredSearch-150M

Benchmarks: LF-WikiSeeAlsoTitles-320K LF-WikiSeeAlso-320K LF-WikiTitles-500K LF-Wikipedia-500K LF-ORCAS-800K SponsoredSearch-150M offline

Metrics: Precision@1 Precision@5 nDCG@5 PSP@1 PSP@5 Recall@100 click-through rate impression yield click yield

Baselines: NGAME ANCE PINA GraphFormers GraphSAGE DEXA DPR

Builds on: NGAME DistilBERT DPR Direct Preference Optimization

Closest papers by meaning

• NUDGE: Lightweight Non-Parametric Fine-Tuning of Embeddings for Retrieval (ICLR 2025 · similarity 0.84)
• xRAG: Extreme Context Compression for Retrieval-augmented Generation with One Token (NeurIPS 2024 · similarity 0.84)
• MOGIC: Metadata-infused Oracle Guidance for Improved Extreme Classification (ICML 2025 · similarity 0.83)
• Select and Augment: Enhanced Dense Retrieval Knowledge Graph Augmentation (Abstract Reprint) (AAAI 2024 · similarity 0.83)
• NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models (ICLR 2025 · similarity 0.83)
• MUVERA: Multi-Vector Retrieval via Fixed Dimensional Encoding (NeurIPS 2024 · similarity 0.82)