Improving Dense Retriever Training with Corrector Networks (2025)

1 Introduction

The softmax function, paired with deep neural encoder models,is often the parameterization of choice for discrete distributionsover many targets such as in classification (Logeswaran etal., 2019; Yu etal., 2022),retrieval (Reddi etal., 2019; Xiong etal., 2020), or reinforcement learning(Dulac-Arnold etal., 2015; Gottipati etal., 2020).This approach, often called a “dual encoder,” employs two separate deep networks, one to map an input to a fixed dimensional vector, another to map targets to the same vector space. We then compute softmax logits as the inner product of an input vector to each target vector (Gillick etal., 2019; Karpukhin etal., 2020; Xiong etal., 2020).

With the typical softmax cross-entropy loss, exact training of the parameters of these two encoder networks would involve using the current parameters to compute the logits for all targets, requiring running the target encoder on all targets at every step of training. Of course, this far-too-burdensome approach is not used in practice. Instead, various approximations have been developed (Reddi etal., 2019; Rawat etal., 2020; Lindgren etal., 2021; Xiong etal., 2020; Monath etal., 2023).The typical approximation computes a truncated softmax on a sampled subset of targets. These approaches store a cache of “stale” encoded representations of targets and uses the stale, cached representations to draw samples from the softmax-parameterized distribution during training (Lindgren etal., 2021; Izacard etal., 2022). Previous work has used these stale representations amidst other approximations such as index structures (Xiong etal., 2020; Monath etal., 2023), kernel-methods (Rawat etal., 2019), and focusing training on subsets of targets (Reddi etal., 2019). However, inevitably, the staleness of the target embeddings causes training regret.

In this work, we present a simple, general purpose method for addressing staleness in softmax-parameterized categorical distributions that is scalable enough to be updated at every step of training. Our approach improves upon an existing stale approximation using a learned target corrector network. The target corrector network, inspired by recent work on training continuous energy-based models (Han etal., 2020; Grathwohl etal., 2020, 2021), is a small parametric model that accounts for the discrepancy between the stale approximation and unnormalized logits from the true distribution. By learning to improve upon the stale approximation, the target corrector network can be used to produce a more accurate approximation to the target distribution. We further extend beyond training large output space classifiers to latent variable retrieval augmented language models.

In summary, the contributions of this paper are:

Methodological (§3) - We describe a novel training procedure for large output space models. It is based on approximating softmax-parameterized categorical distributions by using a parametric target corrector network that learns to improve stale approximations of logits.

Theoretical (§4) - We analyze the generalization properties of the corrector networks in terms of the discrepancy between the stale approximation and the true distribution, the complexity of the network, and the amount of training data.

Empirical (§5) -We evaluate our approach in training both dense retrieval models and latent variable retrieval augmented language models. Our approach matches the performance of much more computationally intensive approaches at a fraction of the computational expense.

2 Background

Top-K / Similarity-based Sampling Approximations  We can instead use an informed strategy using $g(y)$ that would select higher probability targets by sampling using similarity scores via Gumbel-Max (Lindgren etal., 2021), or using the top-k targets in terms of inner product (Xiong etal., 2020).Work has considered efficient approximations to find these top $k$ targets without having to compute $g(y)$ for all $y\in\mathcal{Y}$ (Xiong etal., 2020; Monath etal., 2023).

Initialization  We initialize the parameters of the dual encoders, $\Theta$, using pre-trained models, such as pre-trained language models, T5 and GTR (Devlin etal., 2019; Raffel etal., 2020; Ni etal., 2022).

Stale Cached Representations When we are training the parameters, $\Theta$, each target’s vector according to $g(y;\Theta)$ changes at each step of training. Therefore, a commonly used approach is to define an approximation, $g^{\prime}(y)$, that is a lookup for a “stale” cached embedding for the given target. The stale embedding comes from running the target encoder at a particular time step $t$, of training, and caching the result, i.e., $g(y,\Theta_{t})$ in a buffer, $B\in\mathcal{R}^{|\mathcal{Y}|\times D}$, i.e. $g^{\prime}(y_{i})\triangleq B_{y_{i}}$.To find the top-$k$ targets for input $x$ we compute approximate logits $B^{T}f(x)\in\mathcal{R}^{|\mathcal{Y}|}$ and select the top-$k$ targets to define $S(\mathcal{Y})$. Even before training, we can use the pre-trained model to produce embeddings for all targets $B_{y_{i}}=g(y_{i};\Theta_{0})$. While $B$ may seem large, this is considerably more efficient than exact computation and is possible, on accelerators, for $|\mathcal{Y}|$ in the tens of millions.

The bias of this approach (and subsequent degradation in performance) depends on the staleness or drift of the embeddings, i.e., $||B_{y_{i}}-g(y_{i})||$ which will increase as we update the parameters of $g(y)$. This can be mitigated by recomputing $B$ periodically throughout training (at notable cost). This approach of periodically recomputing has been used (Guu etal., 2020; Izacard etal., 2022; Monath etal., 2023), but there is still much room for improvement.

3 Improving Training with Target Correctors

Our proposed approach builds upon these stale buffer approximations by using an additional parametric model. The additional model aims to improve upon the stale $g^{\prime}(y)$ to yield a better approximation of $g(y)$.

With significantly fewer parameters than a typical dual-encoder, i.e., $|\Psi|\ll|\Theta|$, this small parameteric model is efficient enough to provide approximately fresh representations of every target at every training step. The target corrector network presents interesting research questions regarding whether the network can obviate the need for re-embedding, what kinds of staleness or drift can be effectively modeled, and how much training data is required.

Warmup: Training the corrector network in isolation We begin by considering how we would train only the parameters of the target corrector network, independently of the dual-encoders $f(x)$ and $g(y)$. Afterwards, we present an algorithm for jointly training the target corrector network and the dual-encoders.To train the parameters $\Psi$ of the corrector network, $h(\cdot;\Psi)$, we collect training examples of input data points $x_{i}$, the exact target embeddings $g(y)$, and stale embeddings $g^{\prime}(y)$ for a subset of targets $S(\mathcal{Y})_{i}\subset\mathcal{Y}$, i.e., $\{(f(x_{i}),g(y_{b}),g^{\prime}(y_{b}))\ |\ y_{b}\in S(\mathcal{Y})_{i}\}$.

Jointly Training Corrector Networks & Dual-Encoders We present a method (Algorithm1) for simultaneously training dual-encoders for a given task (e.g., retrieval or equivalently large output-space classification)and the target corrector network. The training algorithm will optimize both theparameters of the target corrector network and additionally use the corrector network to approximate the softmax. Each step consists of: (1) using the corrector network to provide an approximately updated representation of every target, (2) picking a subset of targets for the truncated softmax using the output of the corrector network, (3) computing a task loss for the dual-encoder models and loss for the corrector networks, (4) updating, according to their respective losses, the parameters for both the dual-encoders and the corrector networks using gradient descent.

In more detail, we are given task training data, $X=\{(x_{1},y_{1}),\dots,(x_{m},y_{m})\}$. We are given a task loss function $\mathcal{L}$ and a corrector network loss $\ell$. The dual-encoder models are $f(x),g(y)$ and their initial parameters are $\Theta_{0}$.Prior to the first training step, we instantiate a buffer of the targets’ representations, $B_{y}=g^{\prime}(y)=g(y;\Theta_{0})$. We will avert the need for the expensive updating of the buffer by re-embedding targets with the target encoder.In each step, we sample a training point and label pair $x_{i},y_{i}$ from $X$. We apply the target corrector network to all of the stale representations in the buffer to obtain $h\circ g^{\prime}(y)\ \forall y\in\mathcal{Y}$. This computation does not require running a dual-encoder; we use the cached buffer representation of each target as input to the corrector network. The corrector network is typically a two-layer MLP and hence efficient enough to be used in this way. With these representations from $h(\cdot)$, we sample (or select exact top-$k$) targets according to $P_{h}(y|x)$ (Eq.4) to form a subset of targets $S_{x_{i}}(\mathcal{Y})$ for the truncated softmax.

Given this subset, we compute the task and correction losses and update their respective model parameters. First, we compute the task loss, which is cross-entropy. The task loss will only be used to update the parameters of the dual-encoders, $\Theta$, not the parameters of the target corrector network. We compute the truncated softmax $\tilde{P}(y|x)\propto\exp(\beta\langle f(x),g(y)\rangle)$ (Equation3). We define a one-hot $P^{\star}$ according to the training data label $y_{i}$. We compute the task specific loss $\mathcal{L}$ as a function of $\tilde{P}$ and $P^{\star}$, and update the dual encoder parameters via gradient descent $\Theta\leftarrow\Theta-\eta\nabla_{\Theta}\mathcal{L}$.

Next, we will use the same sample of targets $S_{x_{i}}(\mathcal{Y})$ to compute the target corrector network’s loss and parameter update. Importantly, this will only update the parameters of the target corrector network, $\Psi$, not the parameters of the dual-encoders. Here we describe the use of the cross-entropy loss. However, an analogous update procedure could be used for other loss functions. We compute the truncated softmax according to the target corrector network’s output: $\tilde{P}_{h}(y|x)\propto\exp(\beta\langle f(x),h\circ g^{\prime}(y)\rangle)$. We then compute the target corrector network loss, $\ell$, cross-entropy, which tries to align two truncated distributions $\tilde{P}_{h}$ and $\tilde{P}$. The target corrector network’s parameters are updated by gradient descent $\Psi\leftarrow\Psi-\eta\nabla_{\Psi}\ell$.

Training the target corrector network, which has only a small number of parameters, is much less computationally intensive to train than the dual-encoder model. Furthermore, we are given “for free” the representations $g(y)$ since they are used to compute $\tilde{P}$ for the task loss. These representations can then easily be re-used for training the corrector.

The training procedure is summarized in Algorithm1. At prediction time, the corrector network is not used, instead the trained dual-encoder $g(y,\Theta)$ is used.

4 Analysis

We will explore the generalization of the proposed target corrector network in terms of unseen targets for a particular input data point, and will show the relationship between generalization error, the complexity of the target corrector network $h$, and the discrepancy of the stale representations, $g^{\prime}$, and true representations $g$. All proofs are in AppendixA.

Let $\ell:\mathbb{R}\times\mathbb{R}\to\mathbb{R}$ is a loss function for the target corrector network (Eq.5 & 6).

We begin by defining three kinds of risk.

Function Classes  The function class $\phi\in\mathcal{F}$ is large. We will relate this large function class to a restricted class of functions of the form $h\circ g^{\prime}$ by leveraging the approximate stale representations, $g^{\prime}$.In other words, we restrict $\mathcal{F}$ to $\mathcal{F}^{g^{\prime}}=\{h\circ g^{\prime}:h\in\mathcal{H}\}$ where $\mathcal{H}$ represents the simpler function class mapping $\mathbb{R}^{d}\to\mathbb{R}^{d}$ which can express the discrepancy between the stale $g^{\prime}$ and current $g$.

First, we provide a bound on the gap between the population risk and stale population risk. We formalize this in the following lemma. For ease of notation in this exposition, we define $\mathcal{G}_{\ell,\mathcal{F}}$ as the induced function class:$\mathcal{G}_{\ell,\mathcal{F}}=\{y\mapsto\ell(\phi(y),g(y)):\phi\in\mathcal{F}\}.$

Next, we connect stale population risk to the empirical risk.

Now, we can relate the complexity of function class $\mathcal{F}^{g^{\prime}}$, number of samples $n$, and the discrepancy of the true $g$ and stale approximate encoders $g^{\prime}$:

Note the following implications of these theoretical results:

1. If the corrector network $h$ is too complicated or there are not enough samples $n$, then $h$ overfits and $T_{3}$, will dominate.

2. If $g$ and $g^{\prime}$ are very different, then term $T_{2}$ will dominate.

3. If $h(\cdot)$ is too simple and we cannot fit the sampled data well, then $T_{1}$ will dominate.

We empirically explore some of these trade-offs in §5.3.

5 Experiments

We evaluate training using target corrector networks in two settings: supervised dense retrieval and retrieval augmented language models.We further investigate the properties of the target corrector networks in synthetic experiments.In summary, the experiments investigate whether training strategies can effectively obviate the need to keep cached buffers of targets up-to-date by re-embedding during training. We answer this question affirmatively with the following highlights:

No Re-embedding Needed Training using target corrector networks matches the task performance of exhaustively re-embed all targets every 500 steps throughout training in both dense retrieval (Table1) and retrieval augmented language models (Table3). Target correctors achieve this without ever needing to re-embed targets during training, yielding significant computational savings (Fig.2).

Best no re-embedding method Compared to frozen approaches, stale approaches, and Dynnibal without re-embedding, target corrector networks achieve over 10 point improvements in RLM tasks and 4 point improvements across multiple recall measures in retrieval.

Simpler and Less Computation Target correctors perform as well or better than Stochastic Negative Mining (SNM) (Reddi etal., 2019) despite SNM doing more re-embedding. Similarly, target corrector networks nearly match Dynnibal (Monath etal., 2023) when Dynnibal uses much more computation (Table1). Dynnibal is a much more complicated and difficult to implement method.