Preserving Multilingual Quality While Tuning Query Encoder on English Only

Of ARXIV we use titles and abstracts. In order to have a representative subset of a manageable size for our evaluations, we select all samples that have at least one category with a maximum size of 10K samples. For example, the arxiv category bayes-an is the smallest (size 16) in our snapshot (version 173), meaning that there were only 16 arxiv preprints in this category.

We made two flavors of evaluation arxiv triplets from this arxiv subset. In the first version, the anchor is the title, the positive is the corresponding abstract, and the negative is another random abstract. In the second version the anchor is the first sentence of the ’positive’ abstract, the positive is the rest of the abstract, and the negative is a similar piece (first sentence excluded) of the ’negative’ abstract.

We make use of triplets created from arxiv because this provides our evaluation with a very different kind of text (compared to MSMARCO), and thus allows us to judge robustness of the improvement. For convenience and reproducibility of creating triplets of different levels of difficulty, we made a dataset arxiv-negatives¹⁰¹⁰10https://huggingface.co/datasets/primer-ai/arxiv-negatives. The dataset consists of 253140 samples, each sample is a tuple of two elements:

1. 1.

   (1) An ARXIV paper metadata, including its Id, title and abstract and categories.

2. 2.

   (2) List of 21 Ids of other ARXIV papers. The first 20 Ids are the papers that are ’closest’ to the above paper, and sorted from the most to the least similar; the last 21st Id is an Id of a randomly selected paper (not coinciding with Id of the above paper).

Thus, we have 21 versions of picking up negatives for triplets, from the most difficult to the easiest (the last one, of the random selection).

For example, to create triplets of difficulty 14, for each paper given by the first tuple element, we pick up a paper corresponding to 14th Id given in the second tuple element. From the first paper we can create query and positive, and from the second paper - negative. Through this work we used two flavours:

1. 1.

   ’Title’: Query is the title of the first paper, and positive is its abstract; negative is the abstract of the second paper.

2. 2.

   ’First’: Query is the first sentence of the abstract of the first paper; positive is the rest of the abstract; negative is the abstract of the second paper, with its first sentence deleted.

The above dataset is created from the mirror of arxiv¹¹¹¹11https://huggingface.co/datasets/arxiv-community/arxiv_dataset¹²¹²12https://www.kaggle.com/datasets/Cornell-University/arxiv (version 173) arxiv-metadata-oai-snapshot.jsonl through the following steps:

1. 1.

   Identified all arxiv categories with a maximum size of 10K papers (i.e. arxiv preprints).

2. 2.

   Selected all papers that have at least one of the categories identified above. This is the subset of arxiv to deal with: manageably small, yet diverse.

3. 3.

   For each paper: (1) Sort its categories by size, from smaller to larger. (2) Find all other papers that have the closest match by the categories (the closest match is the longest consecutive list of matched categories, starting from the first one). (3) Of the found papers, select 20 closest by Jensen-Shannon distance between the paragraphs, and sort them by the distance. If there were less than 20 papers, fill to 20 by the last one. (4) Add randomly selected paper as 21st.

Of the total 253140 samples, in 213156 samples (84.2%) all the first 20 negatives are different (which means that not less than 20 papers happen to have the same closest match by categories).

A weight decay may restrict increase of model weights, but it does not improve the evaluation results. We show some representative results in Tables 12 and 13. While restricting gains on the tuning goal, weight decay does not help to preserve the other qualities: the results on XNLI and ARXIV are no better than without weight decay. If there is any recipe for further improving the gains both on the tuning goal and on the related qualities, it has to be a less crude interference into the tuning.

Since weight decay may be more effective at higher learning rates, the parameters for Table 12 are chosen at higher rate and batch size, compared to our ’default’ choice, which is used in Table 13. The learning rates and batch sizes of these tables relate by square root scaling rule (see Section M.1.1).

From evaluation results in Figure 2 we may consider the learning rate below 7e-8 (but above 1e-8) as safely suitable for adiabatic tuning. But we know this only because we evaluated the tuned models on the out-of-tuning domains ARXIV and XNLI.

Is there any way to know the upper boundary without having extensive data for evaluation? Could there be an empirical recommendation not to exceed certain learning rate? Can we increase the adiabatic tuning range of learning rate?

In attempting to answer these questions, we considered the largest changes in the layers at different learning rates. One suspect layer, by simple crude measures, is output.dense.weight. Here in Tables 14 and 15 we show why we selected this layer. Our motivation here is based on a simple and crude criteria; more detailed research and understanding may reveal better ways to extend the adiabatic tuning regime.

The gains from the tuning by freezing the layer output.dense.weight (in each transformer block) are shown in Figure 11. In comparison to the default tuning (Figure 2) we can see that the adiabatic regime indeed extends from a learning rate of about 6e-8 (as was in Figure 2) to about 1.3e-7. Thus, freezing of output.dense.weight did help to somewhat extend the adiabatic tuning regime. However, further increase of the learning rate results in worse deterioration for the version with frozen output.dense.weight layer, as can be seen for XNLI starting from the rate 1.4e-7.

When using the triplet loss for contrastive learning, the margin is an important parameter that can significantly affect model training. In Figure 12 we show the dependency of the evaluation results on the margin during its tuning. We consider our default tuning parameters (Section 2.3), but change the margin. The results are not unexpected: a margin up to 0.15 is reasonable, and at higher margins the disturbance on cross-lingual, and, eventually, on English data evaluation becomes too strong.

The corresponding data for $L12$ are given in Figure 13. It shows that a margin of 0.1 works best for $L12$. The results for margin 0.1 are distinctly better. Altogether, $L12$ appears to be more sensitive (compared to $E5$) to the tuning parameters if the goal is to preserve performance on multilingual XNLI data and on out-of-domain ARXIV data. Arguably, the margin value of approximately 0.1 is the best both for $L12$ and $E5$.

In Table 16 we show how the improvement depends on the stopping criterion. The stoppings after 5 or 10 non-improvement epochs give similar results. Stopping after 15 non-improvement epochs continues the trend of increased gain on English data, but with a deterioration on a few language pairs.

There is no essential difference between the execution times for $E5$ and $L12$. The tuning time depends on how soon stopping happened. At the settings of interest (Section 2.3, 3.1, 3.2), the tuning on an A100 GPU takes about one hour. For example, tuning 10 times at the default settings (Section 2.3, Appendix E) for rates between 1e-8 and 1e-7 takes 9 hours. At higher rates, stopping occurs earlier; tuning 10 times for rates between 1.1e-7 to 2e-7 takes less than 5 hours. Table 1 (with freezing different parts of the encoder) was obtained in 6 hours.

Evaluation of an encoder on all datasets we used (MSMARCO, ARXIV-first, ARXIV-title and XNLI) takes about 1.2-1.3 hours.

Using the fine-tuned $E5$ model with the frozen embedding block, tuned using a batch size of 14, and a margin of 0.1, we randomly vary the batch size, learning rate scheduler and weight decay in order to assess their impact on the model’s final performance. In Table 17 we present the change in performance across these different configurations for a learning rate of $5\times 10^{-8}$, which is our ‘default’ learning rate (Section 2.3). In Table 18 we do the same for a learning rate of $10^{-7}$. Table 19 lists the different schedulers we considered. Values in blue indicate the top improvements whereas values in red indicate the worse degradation.

Across these parameters, on average, the batch size appears to have the most significant impact, generally leading to poorer performance as the batch size is decreased. Within each batch size group, we see that using an exponential learning rate scheduler ($E_{0.95}$ or $E_{0.98}$) is generally worse than using any of the other schedulers or no scheduler at all. A specific exception exists when using a batch size of 100 where the exponential scheduler outperforms the quadratic one when the learning rate is set to $10^{-7}$. Across all the configurations considered, the most impact seems to be the one shown in the first row of Table 17, where we see good improvement over MSMARCO and ARXIV-first while simultaneously showing improvement over XNLI ent-neutr.

Table 20 shows the effects of choosing a different optimizer with a small and large learning rate. In addition to Adamax, we tried Adadelta and Stochastic Gradient Descent (SGD), both of which did not change the model weights in a significant enough way to affect the overall performance and hence, are not presented. For higher learning rates, SGD without momentum did elicit a change as shown in Fig. 14, but the trend in performance is similar to what is presented in Fig. 2 with higher resolution near the transition point between improved and degraded multilingual performance. At around $9\times 10^{-7}$, we see a sharp increase in the number of degraded language pairs while the model maintains constant improvement on MSMARCO. With a high enough learning rate, it seems that the gradients are able to overcome a barrier in the loss landscape that confined the weights to a region in which multilingual characteristics were preserved.

From the table, the default version of Adamax (Adamax with momentum) has a nearly negligible effect on the model when used with a small learning rate, suggesting that this particular configuration for the optimizer forces the model weights to change very slowly. When momentum is switched off, the model weights change enough to improve the overall performance in both English and other languages. Continuing down to the bottom row, if we turn up the learning rate to a higher value, the model weights begin to change more significantly which brings about less improvement in the model’s multilingual capacity (still an improvement nonetheless), but maintains the same improvement on English. Overall, going from the first row to the last row (for Adamax), we transition from a point in model weight space where performance on all languages can be enhanced or preserved to a point which is better suited for the English-only task defined in tuning.