The mixture-of-experts (MoE) architecture has proven indispensable for scaling transformer-based artificial intelligence (AI) models. MoEs use a router network to activate specialized subnetworks, called experts, within each transformer layer on a per-token basis. Thus, MoEs increase model capacity without proportionally increasing the computational requirements of training and inference. However, a key challenge is that MoE inference throughput scales poorly with batching, as each input in a batch activates a different subset of experts, thereby increasing memory intensity.

To improve batched MoE inference throughput, we propose TokenMix, a novel technique that mixes token representations within the router across inputs in a batch to promote expert reuse. In effect, TokenMix encourages coordinated expert selection across inputs. TokenMix requires neither training nor fine-tuning and can be applied directly to off-the-shelf MoE models. We evaluate TokenMix across five state-of-the-art MoE models with parameter counts ranging from 16B to 235B. TokenMix achieves a 25–60% improvement in expert reuse, leading to a 12–28% reduction in inter-token latency with negligible or no loss in accuracy.

Document understanding requires answering questions over long, multi-page documents that interleave text, figures, and tables. Recent methods adopt the Document Retrieval-Augmented Generation (DocRAG) paradigm, which indexes each document and feeds the top-K most relevant pages to a Multimodal Large Language Model (MLLM) per query. Because each page can be represented as extracted text or a rendered image, DocRAG systems expose two key knobs: an indexing modality choice when the documents are ingested and a context modality at query time. Existing systems either fix both knobs to the visual modality or use both modalities throughout, incurring indexing and query costs regardless of document or query type. We propose AdaptiveDocRAG, which adapts both knobs as well as the size of visual context to optimize the trade-off between accuracy, query response time, and indexing time. AdaptiveDocRAG achieves this through three lightweight modules. First, document-adaptive indexing modality selection uses a text-density heuristic to route text-heavy documents to cheaper textual indexing. Second, query-adaptive context modality selection uses a learned per-query router to select textual or visual context. Third, query-aware visual context compression reuses the retriever’s per-patch query-page similarity to crop relevant sub-page regions and fuse them with top-ranked text chunks. On MMLongBench-Doc and NQ-Text with Qwen3-VL-8B, AdaptiveDocRAG matches or improves visual-only DocRAG accuracy (up to +1.1 points) at ~50% of the MLLM tokens, with 1.30-1.79x end-to-end speedup and 1.34x faster indexing.

The great success of Language Models (LMs) for various Natural Language Processing (NLP) tasks is accompanied by computational challenges during both pre-training and fine-tuning. Pre-training has attracted significant attention due to its huge computational footprint. We focus on the fine-tuning of pre-trained LMs, which is expected to be performed much more frequently as the pre-trained models are adapted to downstream tasks. During fine-tuning, the presence of variable-length input sequences necessitates the use of padding tokens when batching sequences. These padding tokens lead to ineffectual computations, adversely impacting the efficiency of fine-tuning. We also observe that LMs memorize the limited task-specific training data despite the use of known regularization methods. Based on these insights, we present TokenDrop + BucketSampler, a framework that simultaneously improves efficiency and accuracy of LM fine-tuning. BucketSampler generates batches of samples with lower variance in sequence lengths to reduce the number of padding tokens, but does so without the accompanying accuracy drop seen in previous approaches. TokenDrop is a new regularizer that prunes a random subset of insignificant tokens from each input sequence in every epoch to prevent overfitting. TokenDrop drops more tokens from the longer sequences in each batch to further reduce variance in input lengths and the need for padding. TokenDrop + BucketSampler accelerates fine-tuning on diverse downstream tasks by up to 10.61X, while also producing models that are up to 1.17% more accurate compared to conventional fine-tuning. Code is available at https://github.com/amrnag/TokenDrop-BucketSampler. .