ARMOR

Adaptive Regularized Mixture Optimization for Retrievers

Heshan Fernando Quan Xiao Yan Xin Tianyi Chen

Abstract

Telecom question answering (QA) is a challenging setting for retrieval-augmented generation (RAG): evidence is fragmented across standards, papers, encyclopedic resources, and web documents, and answers often hinge on technical tables, equations, and specialized protocol language. In low-resource subdomains, generator fine-tuning can over-specialize and degrade general capability, making query-side retriever adaptation an attractive alternative.

To this end, we ask whether a fixed-generator, query-adapted RAG system can outperform generator-side adaptation, and which retriever objectives best support that setting. We argue that retrieval, not model fine-tuning, is an effective adaptation target: theoretically, we show that when domain data are scarce, query-encoder tuning can have lower estimation complexity than supervised fine-tuning under standard capacity assumptions.

We identify two particularly relevant objectives—the latent-document RAG likelihood, which optimizes generation utility, and the InfoNCE contrastive objective, which improves semantic retrieval geometry—and leverage them jointly through a retriever optimization method targeting downstream QA performance in the telecom domain. Specifically, we introduce ARMOR (Adaptive Regularized Mixture Optimization for Retrievers), which learns separate temperatures for the RAG retrieval distribution and InfoNCE softmax and regularizes the adapted query encoder toward the frozen base query encoder. Across telecom-specific retrieval and generative QA benchmarks, we show that ARMOR improves evidence retrieval and answer generation in several in-domain settings.

Methodology

Retriever-Centric Adaptation

In the ARMOR setup, documents are embedded once using a base dense retriever and stored in a fixed index. During adaptation, only the query encoder is updated. This focuses limited supervision on the component that controls which evidence is shown to the model, ensuring the document representation remains stable.

By bypassing costly document re-indexing and avoiding full-scale generator fine-tuning, query-side adaptation provides a computationally efficient path for domain specialization, particularly in low-resource regimes.

Retriever adaptation motivation
Retriever-side query-encoder adaptation provides a strong low-resource adaptation path compared with generator-side tuning and other baselines.

Adaptive Regularized Mixture Optimization

ARMOR optimizes a combination of three key components to adapt the retriever toward downstream generation utility, balancing retrieval quality and model stability:

1. RAG Likelihood Objective
fRAG(x, y; τr) = -log ∑d ∈ Dk(x) pηr)(d|x) pθ0(y|x, d)

Rewards evidence that helps the frozen generator (pθ0) produce the correct answer. The learned temperature τr controls how sharply the retrieval probability distribution focuses on high-utility passages.

2. Contrastive InfoNCE Objective
fNCE(x, d+, N-; τc) = -log [ exp(sη(x, d+)/τc) / (exp(sη(x, d+)/τc) + ∑d- ∈ N- exp(sη(x, d-)/τc)) ]

Pulls query-positive pairs (d+) together and pushes negative distractors (d-) apart to stabilize retrieval geometry. The learned temperature τc adaptively scales contrastive margins.

3. Query Distillation Regularizer
fqdist = 1Bb=1B ( 1 - cos(qη(xb), q0(xb)) )

Anchors the adapted query encoder embeddings (qη) to the frozen base query encoder (q0). This prevents query drift and preserves indexing compatibility.

ARMOR Unified Objective
fARMOR = fRAG(x, y; τr) + fNCE(x, d+, N-; τc) + λq fqdist

Rather than fixing a static mixture weight, the optimizer dynamically learns the temperatures τr and τc to balance objectives across different training stages.

Key Results

In-Domain Performance (Tele-Eval)

ARMOR provides the strongest overall performance across answer quality (Score) and retrieval recall (Recall@3 / Recall@5) across the ISAC, JCC, and SAGIN splits compared to other baselines.

Method ISAC JCC SAGIN
Score R@3 / R@5 Score R@3 / R@5 Score R@3 / R@5
Base Gen 0.227 -- 0.298 -- 0.302 --
Base RAG 0.689 0.740 / 0.807 0.776 0.613 / 0.700 0.766 0.813 / 0.840
RAG QE FT 0.658 0.653 / 0.740 0.723 0.600 / 0.653 0.757 0.693 / 0.787
InfoNCE QE FT 0.669 0.673 / 0.753 0.743 0.620 / 0.680 0.766 0.740 / 0.807
Mix QE FT 0.685 0.633 / 0.713 0.736 0.640 / 0.680 0.759 0.720 / 0.780
ARMOR (Ours) 0.712 0.767 / 0.820 0.772 0.647 / 0.713 0.769 0.840 / 0.847

Component Ablation: Coupling Adaptivity and Regularization

Evaluating individual components of ARMOR reveals that adaptive temperatures and query distillation regularization are deeply coupled; omitting either drops performance below static mixture baselines.

Method Adaptive Temp. Regularization Score (ISAC)
Base RAG -- -- 0.689
RAG QE FT -- -- 0.658
InfoNCE QE FT -- -- 0.669
Mix QE FT No No 0.685
Static Mix with Reg. No Yes 0.673
Dynamic Mix without Reg. Yes No 0.635
ARMOR Yes Yes 0.712

Robustness Across Generator Scale

We evaluate Base Gen, Base RAG, and ARMOR across generator backbones of varying scales (Llama-3-8B-Instruct, Llama-3.2-3B, Llama-3.2-1B, and Qwen3-8B). While 1B and 3B generators show limited headroom, 8B generators leverage the improved evidence selection of ARMOR more effectively. ARMOR gains are consistent across 8B model families, demonstrating that retriever-side adaptation becomes increasingly valuable as the generator capacity improves.

ARMOR results across generator backbones

Adaptive Training Dynamics

The retrieval temperature sharpens during training (as τr decreases), showing that the retriever increasingly focuses on high-utility evidence. Query distillation helps constrain this adaptation (measured via query distillation loss) so the tuned query encoder remains compatible with the frozen document embedding space. Select a domain below to view its specific training trajectories:

ISAC training dynamics
ISAC Adaptive retriever temperature (τr, τc) and query distillation loss dynamics.
JCC training dynamics
JCC Adaptive retriever temperature (τr, τc) and query distillation loss dynamics.
SAGIN training dynamics
SAGIN Adaptive retriever temperature (τr, τc) and query distillation loss dynamics.

Data and Corpus Ablations

Evaluating retriever adaptation under different data constraints reveals that training data composition is just as crucial as size in low-resource regimes.

Data size ablation
Training Data Size: Varying the number of training QA pairs shows that gains are positive but non-monotone under extreme low-resource conditions.
Document count ablation
Corpus Coverage: Varying the number of source documents tests coverage under the source-document split constraints.
Document type ablation
Document Type: Restricting training to specific document types (standards, arXiv, Wikipedia) shows how evidence source impacts downstream QA quality.

Discussion & Takeaways

Key insights from the interface-layer retriever adaptation study under real training and deployment constraints:

Which component moves?
During adaptation, only the query encoder is updated. The generator and the document index/encoder remain completely frozen. This ensures we do not have to rebuild the massive vector index when training on new domains.
Which objectives matter?
We couple two key views of RAG optimization: the RAG likelihood for capturing answer generation utility, and the contrastive InfoNCE objective for stabilizing and shaping semantic retrieval geometry.
Why temperature adaptivity?
Instead of a fixed trade-off weight, ARMOR learns separate temperatures (τr, τc). This allows the optimization to dynamically shift emphasis as training progresses from coarse evidence search to fine retrieval sharpening.
Why query regularization?
Unconstrained query-side adaptation can easily drift away from the frozen document space. Query distillation anchors the adapted encoder to the base encoder, preserving index compatibility while allowing domain alignment.

Try ARMOR-powered Telecom Agent

Domain Selection

Select domains to restrict the retriever, or leave empty for automatic routing.