Large Language Models Meet Extreme Multi-label Classification: Scaling and Multi-modal Framework

AI summary

Overview

Research area: Extreme Multi-label Classification (XMC), combining decoder-only large language models and multi-modal (vision + language) learning, with evaluation on Amazon product-to-product recommendation benchmarks.

Technical level: Intermediate. Readers need some familiarity with transformer architectures, contrastive learning, dual-encoder retrieval, and standard XMC metrics (P@k, PSP@k, R@k), but the paper is structured around two intuitive ideas: scale up the text backbone, and add one image embedding per item.

One-sentence scope: The paper shows that decoder-only LLMs can be adapted into Siamese-style XMC pipelines (dual-decoder learning) and that efficiently injecting frozen vision-model image embeddings (ViXML) improves XMC accuracy enough that a 66M-parameter encoder-based model outperforms billion-parameter text-only models in most cases.

What This Paper Is About

XMC asks a system to pick the most relevant labels for a query out of an extremely large label space, often containing millions of labels, while staying computationally efficient. Most recent methods treat this as a maximum inner product search between query and label embeddings learned by small encoder-only transformers, and the paper argues that decoder-only LLMs and visual metadata have both been largely untapped. The goal is to exploit larger decoder-only models and image metadata without letting computational cost grow out of control.

Key Contributions

1. Dual-decoder learning for XMC. A method for adapting decoder-only transformer architectures to Siamese-style XMC by embedding queries and labels inside structured prompt templates and extracting sentence embeddings, scaled up to 7B parameters.
2. ViXML, a vision-enhanced multi-modal framework. An architecture-agnostic framework that integrates a single embedding per image from a frozen foundation vision model into any Siamese-style XMC method, supporting both encoder and decoder backbones while limiting computational growth.
3. Three extended public datasets. The authors add visual metadata from Amazon Reviews to LF-AmazonTitles-131K, LF-Amazon-131K and LF-AmazonTitles-1.3M, and make them publicly available for multi-modal XMC benchmarking.
4. Systematic scaling and comparative study. Experiments across four public datasets, multiple LLM families, and several XMC methods (NGAME, DEXA, PRIME) that quantify the effect of backbone size and vision backbones, reaching improvements of +5.07% to +8.21% points in P@1 over the best previous results.

Main Findings

• Scaling to decoder-only LLMs improves text-only XMC. On LF-AmazonTitles-131K, text-only P@1 rises from 41.72 (MiniLM-L3) and 44.86 (DistilBERT) to 45.87 (BERT), 47.42 (Qwen2.5-3B-Instruct) and 48.06 (Qwen2.5-7B-Instruct). Similar encoder-to-decoder gains appear on MM-AmazonTitles-300K (52.95 BERT vs 54.71 Qwen2.5-3B-Instruct), LF-AmazonTitles-1.3M (59.10 BERT vs 60.74 Qwen2.5-3B-Instruct) and LF-Amazon-131K (49.15 BERT vs 53.17 Qwen2.5-3B-Instruct).
• Decoder models are more sample efficient. Encoders are trained for 300 epochs while decoders are trained for 30, reducing training epochs by an order of magnitude; the 0.5B decoder and the 66M DistilBERT train in comparable time in the authors' setup.
• A 66M encoder with ViXML beats billion-parameter text-only models in most cases. ViXML with DistilBERT reaches 49.55 P@1 on LF-AmazonTitles-131K versus 48.06 for the text-only 7B Qwen2.5-Instruct, and 64.17 on LF-AmazonTitles-1.3M versus 60.74 for the text-only 3B model. The authors summarize this as "an image is worth billions of parameters."
• One exception is reported. Encoder-based ViXML does not surpass the text-only 3B model on LF-Amazon-131K. The authors argue decoder-only models benefit from the longer textual inputs (titles plus descriptions) in that dataset.
• Visual metadata helps every XMC method tested. Adding ViXML on LF-AmazonTitles-131K raises P@1 from 42.47 to 47.81 with NGAME, from 44.38 to 49.48 with DEXA, and from 44.86 to 49.55 with PRIME.
• Early fusion beats late fusion. On MM-AmazonTitles-300K, ViXML with SigLIP2 reaches 55.03 P@1 versus 53.44 for MUFIN-style late fusion on top of PRIME, and ViXML with ViT-32 reaches 53.30 versus 52.62 for MUFIN with ViT-32. ViXML is also reported at +1% P@1 over the original MUFIN (53.30 vs 52.30) while requiring only encoder training.
• Bigger vision backbones help. Performance improves when scaling from the 86M ViT-32 to the 1.14B SigLIP2 in both MUFIN and ViXML.
• Off-the-shelf embeddings underperform. On LF-AmazonTitles-131K, the general-purpose Qwen3-Embedding-4B model reaches 22.33 P@1, and the instruction-prompted variant (marked with an asterisk) reaches 18.33, both far below trained dual-decoder alternatives such as 47.42 (Qwen2.5-I, 3.09B), 47.51 (Qwen3-Emb, 4.02B), 47.74 (Gemma-3-I, 4.3B) and 48.19 (Llama-3.2, 3.21B). The pre-trained vision-language embedding model MoCa (3.75B) reaches 25.67 P@1.
• Prompt structure and image placement matter. In the prompting study on LF-AmazonTitles-131K with Qwen2.5-0.5B-Instruct trained for 20 epochs, P@1 goes from 43.14 with plain text embeddings to 44.15 with a text prefix, 44.66 adding an end-of-sequence token, 49.45 placing images after text, and 49.67 with the adopted template 𝒯1 ⊕ ℰi ⊕ ℐ2 ⊕ 𝒱i ⊕ eEOS. Putting image tokens first (𝒱 ⊕ ℰ ⊕ eEOS) gives 46.12, which the authors attribute to breaking pretraining dynamics and attention sinks.
• ViXML's overhead is bounded. For decoder models, ViXML adds roughly 15%–17% overhead, mostly from longer inputs due to prompts and visual metadata, while encoder-based ViXML halves training epochs from 300 to 150 thanks to faster convergence.
• Full comparison against prior work. Final P@1 for ViXML is 53.08 on LF-AmazonTitles-131K (best prior: MOGIC 47.01), 57.37 on MM-AmazonTitles-300K (MUFIN 52.30), 67.83 on LF-AmazonTitles-1.3M (PRIME 59.62), and 55.57 on LF-Amazon-131K (MOGIC 50.05). The paper notes MOGIC reports +0.88% in PSP@1 relative to ViXML on LF-AmazonTitles-1.3M while dropping 16.88% in P@1.

Methodology in Plain English

The researchers treat XMC as a retrieval problem: turn every query and every label into a single embedding vector, then predict the labels whose vectors have the highest inner product with the query vector. Training uses contrastive learning — a triplet loss (Equation 2) with a margin, L2-normalized embeddings, and NGAME hard negative mining for efficiency — so positive query-label pairs are pulled together and negatives pushed apart.

For the text-only "dual-decoder" variant, each query and each label is wrapped in a short structured prompt made of a text prefix ("This product text"), the text itself, and an end-of-sequence token ("<|endoftext|>"). The decoder LLM is fine-tuned with LoRA and the sentence embedding is obtained by mean pooling. Keeping the prompts short constrains memory.

For ViXML, images are added. A frozen pre-trained vision model maps each image to an embedding, which is then projected by a single learnable linear layer into the text embedding space. Because the vision model is frozen, image embeddings can be stored as a feature bank, and because only one embedding per image is used, sequence length barely grows. With encoder backbones, image embeddings are simply concatenated in front of the text token embeddings (Equation 3). With decoder backbones, the images are placed after the text inside a template that also includes an image prefix such as "and its image" (Equation 4), which the authors found necessary for robust performance.

Experiments use MiniLM-L3, DistilBERT and BERT as encoders; Qwen2.5-Instruct (0.5B, 3B, 7B) plus Llama-3.2, Gemma-3, Qwen3 and Qwen3-Embedding as decoders; and SigLIPv2 (with ViT-32 in the fusion comparison) as the vision encoder. Up to 3 images are used per item, and everything is trained on a single 80GB GPU (an H100 per the supplementary material). Performance is always reported at the last training epoch, with no checkpoint selection on validation metrics.

Why This Matters

The paper connects two research threads that are usually separate — LLM scaling for embeddings and multi-modal XMC — and shows the combination yields larger gains than either alone at a fraction of the expected cost. It also questions the assumption that billion-parameter text-only models are the best route to XMC performance, since a 66M encoder plus frozen image features surpasses them in most evaluated settings. For practitioners, the practical message is that frozen foundation vision models plus a single linear projection and one embedding per image can be a cheap, high-value upgrade to an existing Siamese XMC pipeline.

Real-world applications:

• Product-to-product recommendation, the exact setting of the Amazon-derived benchmarks used here.
• Sponsored search and ad retrieval, where queries must be matched against extremely large candidate label sets.
• Document tagging and topic assignment over large taxonomies, where category or hyperlink metadata has already been shown to help.
• Any catalogue scenario with images, such as marketplaces, where visual item metadata is available at both training and inference time.

Industry relevance: the frozen vision encoder design and the single-embedding-per-image choice mean image features can be precomputed once and reused, avoiding the cost of fine-tuning large VLMs or appending hundreds of visual tokens. The paper reports the whole experimental setup fits on one 80GB GPU, and that the code is available at https://github.com/DiegoOrtego/vixml.

Future Directions

• Attention direction. The authors keep uni-directional attention for decoder models and explicitly leave open whether bi-directional or uni-directional attention works best.
• Richer visual adaptation. ViXML uses a single linear layer to adapt image embeddings to the text space; the paper says studying more complex adaptation choices is left for future work.
• Metadata-free inference. The paper situates itself against prior work (OAK, MOGIC) that assumes metadata is available at training but not test time, leaving open how ViXML's inference-time reliance on images trades off in deployment settings without images.
• Cost of decoder inference. The authors acknowledge that inference throughput will always be impacted by larger decoder parameter sizes and point to specialized hardware or kernel-based implementations such as vLLM as the path to making small LLMs practical for XMC, without evaluating that here.

Target Audience

Researchers and engineers working on extreme multi-label classification, large-scale retrieval, or embedding-based recommendation who want to know whether scaling to decoder-only LLMs and adding visual metadata is worth the cost. It is also useful for practitioners building multi-modal retrieval or product recommendation systems who need a concrete, compute-conscious recipe (frozen vision encoder, one embedding per image, LoRA fine-tuning, single 80GB GPU), and for those who need the three newly released image-enhanced Amazon benchmark datasets for comparison.

Authors’ abstract

Foundation models have revolutionized artificial intelligence across numerous domains, yet their transformative potential remains largely untapped in Extreme Multi-label Classification (XMC). Queries in XMC are associated with relevant labels from extremely large label spaces, where it is critical to strike a balance between efficiency and performance. Therefore, many recent approaches efficiently pose XMC as a maximum inner product search between embeddings learned from small encoder-only transformer architectures. In this paper, we address two important aspects in XMC: how to effectively harness larger decoder-only models, and how to exploit visual information while maintaining computational efficiency. We demonstrate that both play a critical role in XMC separately and can be combined for improved performance. We show that a few billion-size decoder can deliver substantial improvements while keeping computational overhead manageable. Furthermore, our Vision-enhanced eXtreme Multi-label Learning framework (ViXML) efficiently integrates foundation vision models by pooling a single embedding per image. This limits computational growth while unlocking multi-modal capabilities. Remarkably, ViXML with small encoders outperforms text-only decoder in most cases, showing that an image is worth billions of parameters. Finally, we present an extension of existing text-only datasets to exploit visual metadata and make them available for future benchmarking. Comprehensive experiments across four public text-only datasets and their corresponding image enhanced versions validate our proposals' effectiveness, surpassing previous state-of-the-art by up to +8.21\% in P@1 on the largest dataset. ViXML's code is available at https://github.com/DiegoOrtego/vixml.