Improving Dual-Encoder Training through Dynamic Indexes for Negative Mining
This work addresses a key bottleneck in training dual-encoder models for classification and retrieval, offering significant efficiency gains for large-scale applications.
The paper tackled the inefficiency of static nearest neighbor indexes in dual-encoder training by introducing a dynamic tree-based algorithm and a low-rank approximation for target encodings, resulting in cutting error by half compared to brute-force negative mining and using 150x less accelerator memory while surpassing prior state-of-the-art.
Dual encoder models are ubiquitous in modern classification and retrieval. Crucial for training such dual encoders is an accurate estimation of gradients from the partition function of the softmax over the large output space; this requires finding negative targets that contribute most significantly ("hard negatives"). Since dual encoder model parameters change during training, the use of traditional static nearest neighbor indexes can be sub-optimal. These static indexes (1) periodically require expensive re-building of the index, which in turn requires (2) expensive re-encoding of all targets using updated model parameters. This paper addresses both of these challenges. First, we introduce an algorithm that uses a tree structure to approximate the softmax with provable bounds and that dynamically maintains the tree. Second, we approximate the effect of a gradient update on target encodings with an efficient Nystrom low-rank approximation. In our empirical study on datasets with over twenty million targets, our approach cuts error by half in relation to oracle brute-force negative mining. Furthermore, our method surpasses prior state-of-the-art while using 150x less accelerator memory.