MoPo: Momentum Posterior Regularization

• MoPo is a multi-hop dense retrieval framework that uses query-focused hop-wise summaries and momentum updates to stabilize knowledge distillation.
• It replaces traditional answer-based supervision with intermediate, query-specific summaries, reducing semantic drift during multi-hop retrieval.
• Empirical evaluations on HotpotQA and StrategyQA demonstrate substantial gains in recall and exact match without increasing inference costs.

Momentum Posterior Regularization (MoPo) is a framework for multi-hop dense retrieval in open-domain question answering (QA) that enables stable, effective knowledge distillation from a posterior retriever possessing oracle information into a practical prior retriever used at inference time. MoPo addresses the key challenges of posterior regularization in multi-hop settings by (1) replacing answer-based supervision with hop-wise, query-focused summaries and (2) introducing a momentum-based parameter update that constrains the teacher–student knowledge gap during optimization. Extensive empirical results on HotpotQA and StrategyQA demonstrate substantial performance improvements over competitive baselines without increasing inference cost (Xia et al., 2024).

1. Problem Setting and Model Architecture

In multi-hop dense retrieval, given an open-domain question $q$ and large corpus $\mathcal{D}$, the task is to retrieve a sequence of $L$ passages $D_\text{seq} = \{d_1, \ldots, d_L\}$ such that all requisite knowledge for answering $q$ is gathered. The joint probability is factorized as

$$P_\theta(D_\text{seq}|q) = \prod_{t=1}^L P_\theta(d_t | q, d_1, \ldots, d_{t-1})$$

At each hop $t$, the query is reformulated as $q_t = G_s(q_{t-1}, d_{t-1})$, where $G_s$ is a fixed function that concatenates the original $q$ with a short summary $s_{t-1}$ describing retrieved passages up to hop $t-1$. The prior retriever $M_\theta$ (parameters $\theta$) is deployed at inference with access only to $q_t$, whereas the posterior retriever $M_\varphi$ (parameters $\varphi$) is used in training and is privileged with access to query-focused summary $s_t$ reflecting gold knowledge of both previous and current hops.

2. Innovations in Posterior Regularization

2.1 Hop-wise Query-Focused Summaries

Conventional one-hop posterior regularization uses the final answer as posterior information for the teacher model. In multi-hop retrieval, the final answer is often decoupled from intermediate hops, rendering this approach ineffective. MoPo introduces hop-wise query-focused summaries $s_t$ that (i) fuse the current gold passage $d_t$ with preceding contextual summaries, and (ii) explicitly condition on the original question $q$. These summaries function as compact, semantically anchored proxies for the gold context required to inform each retrieval step, and help avoid semantic drift typical in vanilla document concatenation.

2.2 Momentum-Based Posterior Updates

Standard knowledge distillation schedules pretrain the posterior model to convergence, then distill to the prior. This results in significant divergence ("knowledge gap") between $M_\theta$ and $M_\varphi$, causing unstable KL gradients or even performance degradation. MoPo introduces a momentum update for the posterior model:

$$\varphi^{(t)} \leftarrow m \cdot \varphi^{(t-1)} + (1-m) \cdot \theta^{(t-1)}$$

where $0 < m < 1$ is the momentum coefficient. This exponential moving average maintains proximity between posterior and prior parameters throughout training, enabling a stable, smoothly varying KL regularization signal.

3. Mathematical Formalization

Both $M_\theta$ and $M_\varphi$ are dual-encoder models mapping queries $x$ (either $q_t$ or $s_t$) and documents $d$ into vector embeddings $(e_x, e_d)$ scored by dot-product $f_\theta(x,d) = e_x^\top e_d$. At each hop $t$:

• The prior defines $$p_\theta(d_t|q_t) = \frac{\exp(f_\theta(q_t,d_t))}{\sum_{d \in \{d_t^+ \cup d_t^-\}} \exp(f_\theta(q_t,d))}$$
• The posterior analogously defines $p_\varphi(d_t|s_t)$

The overall objective for a batch $\mathcal{B}'$ of $(q, S_\text{seq}, D_\text{seq})$ triples is

$$\mathcal{L}(\theta) = \mathcal{L}_\text{InfoNCE}(\theta;\mathcal{B}') + \lambda\, \mathbb{E}_{(q,S_\text{seq},D_\text{seq}) \in \mathcal{B}'} \Big[ \sum_{t=1}^L D_\text{KL}\left( p_\varphi(\cdot|s_t) \Vert p_\theta(\cdot|q_t) \right) \Big]$$

where $\mathcal{L}_\text{InfoNCE}$ is the standard multi-hop contrastive loss using in-batch negatives. Only $\theta$ receives direct gradient updates; $\varphi$ is updated via the momentum rule above.

4. Training Procedure and Implementation

Training proceeds as follows: initialize $\varphi = \theta$; for each batch, at each hop, (a) build $q_t$ as $q \oplus s_{t-1}$; (b) compute prior and posterior logits over candidate passages; (c) evaluate $\mathcal{L}_\text{InfoNCE}$ and the hop-wise KL term; (d) update $\theta$ via Adam and $\varphi$ by momentum. Passage negatives $d_t^-$ are sampled in-batch from other positives or randomly from the corpus. The loss is summed over all hops, and experiment results show that momentum values $m$ in $[0.9, 0.99]$ are optimal; lower $m$ increases the risk of destabilizing the exponential averaging.

5. Empirical Evaluation

MoPo is evaluated on HotpotQA and StrategyQA for multi-hop retrieval and downstream QA performance.

• Retrieval: On HotpotQA, MoPo achieves recall@2 $94.77\%$ and EM@2 $63.03\%$, outperforming iterative DPR (MDR) (recall@2 $94.34\%$, EM@2 $55.96\%$). On StrategyQA, MoPo attains recall@2 $43.36\%$ and EM@2 $31.91\%$ compared to MDR's $42.64\%$ and $25.31\%$.
• Posterior regularization baselines: Two-stage PR_fixed and PR_dyn underperform both the MDR_sum (concatenation+summarization) baseline and MoPo, confirming that naïve KL regularization is unstable for multi-hop QA.
• Ablations: MoPo’s performance is robust to $\lambda$ (performance gap $\leq3$‰ vs. $>7$‰ for PR_fixed) and is best with momentum $m$ in $[0.9, 0.99]$.
• Downstream QA and reranking: With a lightweight cross-encoder reranker, MoPo achieves EM@2 $89.4\%$ on HotpotQA full-wiki, outperforming BeamDR and Chain-of-Skills. In generation (retrieval→rerank→Flan-T5), MoPo ties prior SOTA joint EM $45.7\%$ and exceeds original MDR ($41.8\%$). On 100-sample Hotpot, MoPo (F1 $64.3\%$) outperforms BeamAggR (F1 $62.9\%$) (Xia et al., 2024).

6. Discussion, Insights, and Limitations

The use of momentum to couple the posterior model to a moving average of the prior prevents divergence between the models (“runaway” teacher), keeps the KL teacher signal consistent, and yields smooth, rapid convergence. Query-focused summaries as posterior information provide strong semantic anchors, reducing drift and improving retrieval accuracy at each hop. MoPo generalizes effectively: training on HotpotQA summaries alone increases StrategyQA EM by more than $6\%$ over the MDR baseline. A plausible implication is that this hop-wise distillation may facilitate transfer to new domains with similar reasoning structures.

Persisting limitations include the absence of a complete theoretical analysis of how momentum regularization constrains the teacher–student divergence, and incomplete benchmarking against LLM-based multi-hop retrievers at inference. Further work could clarify the relationship between the smoothing hyperparameter $m$ and convergence stability, and evaluate potential integration of MoPo with retrieval-augmented generation frameworks.

7. Summary Table: MoPo Algorithm Components

Component	Role in MoPo	Notes on Contrast/Baselines
Query-focused summaries	Compact, hop-wise posterior info at train time	Avoids limitations of answer-based
Momentum posterior update	Couple posterior (φ) to prior (θ) via EMA	Prevents “student-teacher drift”
Dual-encoder scoring	Dot-product retrieval over queries and documents	Standard in dense retrieval
InfoNCE + KL regularization	One-stage loss with hop-aggregated KL	Unstable in two-stage (non-MoPo) PR

The MoPo framework provides a robust and scalable recipe for regularized multi-hop retrieval and QA through hop-wise posterior summarization and momentum-constrained distillation, offering measurable benefits on multi-hop benchmarks without added inference overhead (Xia et al., 2024).