Unleash Model Capacity for Universal Dense Retrieval by Task Specialty Optimization
Authors: Wenzheng Zhang
Chenyan Xiong
Karl Stratos
Arnold Overwijk
Venue: Submitted to ICLR 2023 (2023)
Keywords: Dense Retrieval
Multi-task
Parameter sensitivity

Abstract: Universal dense retrieval, with one unified representation space to empower various retrieval scenarios, has many appealing advantages in simplicity, efficiency, and potential to break echo chambers with cross-scenario information access. However, standard multi-task trained dense retrievers often fail to meet the accuracy of scenario-specific models. In this paper, we analyze the multi-task learning in universal retrieval and show that the model capacity is not the main bottleneck. It is the optimization failed to fully utilize the network parameters to capture task-specific signals. This motivated our development of TACO-DR, which conducts multi-task learning for universal retrieval with TAsk speCialty Optimization. TACO-DR dynamically adjusts the learning rate for each parameter regrading each task based on its task-specific sensitivity, to encourage parameters to better capture task specific signals. On the KILT benchmark, TACO-DR outperforms various multi-task learning methods and achieves better overall accuracy than single-task models. Our analysis shows that TACO-DR better utilizes the model capacity with more task-specific parameters. Our code and model checkpoints will be open-sourced.

Under review as a conference paper at ICLR 2023
UNLEASH MODEL CAPACITY FOR UNIVERSAL DENSE RETRIEVAL BY TASK SPECIALTY OPTIMIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
Universal dense retrieval, with one unified representation space to empower various retrieval scenarios, has many appealing advantages in simplicity, efficiency, and potential to break echo chambers with cross-scenario information access. How- ever, standard multi-task trained dense retrievers often fail meet the accuracy of scenario-specifer models. In this paper, we analyze the multi-task learning in uni- versal retrieval and show that the model capacity not the main bottleneck, is the optimization failed to fully utilize the network parameters to capture task-specific signals. This motivated our development of TACO-DR, which conducts multi-task learning for universal retrieval with TAsk speCialty Optimization. TACO-DR dynamically adjusts the learning rate for each parameter regrading each task based its task-specific sensitivity, to encourage parameters better capture task spe- cific signals. On the KILT benchmark, TACO-DR outperforms various multi-task learning methods and achieves better overall accuracy than single-task models. Our analysis shows that TACO-DR better utilizes the model capacity with more task-specific parameters. Our code and model checkpoints will be open-sourced.
INTRODUCTION
With pretrained language models (Lee et al.) 2019) and dedicated training strategies (Karpukhin et al. 2020a; Xiong et al. 2021), I, dense retrieval systems now effectively learn a dense representation space that matches queries and relevant documents in nearest neighborhoods. This representation--aase retrieval approach provides strong empirical benefits various scenarios with retrieval as the end
goal BBaaa al. 2016) and as the first stage retrieval of many language systems (Lewis et al. 2020).
promising potential of dense retrieval to unify various scenarios via one representation space, that unifies the representation and match of different types of information, e.g., text and image (Liu et al. 2022). and different types of queries, e.g., keywords, questions, and conversations (Petroni al. 2021). Such universal retrieval (Maillard et al. 2021 leads to instant efficiency benefits, as one document index can support multiple scenarios. It also helps break information boundaries between scenarios with one unified entrance for all user information needs.
Ideally, universal retrieval on multiple scenarios would lead to more accurate retrieval than single scenario systems, with the advantage of multi-task learning. However, recent research observed ambivalent empirical performance of universal retrieval, especially when capturing a large number of retrieval tasks one universal system (Maillard et al., (022). This clouds the promise of universal retrieval as its becomes a trade-off between efficiency and accuracy.
In this paper, we conduct thorough investigation on the challenges of multi task learning in universal retrieval. We performed analysis on the KILT benchmark and found that several state-of-the-an retrieval systems indicate that the network capacity is not yet the main limiting factor for universal retrieval accuracy. Though the multi-task learning has guided the parameters capture task specific or shared signals, the optimization not sufficient, resulting in a large fraction of parameters that not well utilized to capture task-specific signals, reflected by their low sensitivity (Liang et al. 2022] to each task.
Motivated by observations, develop TACO-DR, "TAsk speCific Optimized universal Dense Retriever", that improves universal retrieval by optimizing the task-specialty of neural parameters during multi-task training. TACO-DR first utilizes task identifier prompts in its query encoder to

---

Under review as a conference paper at ICLR 2023
improve the model's task awareness without compromising on the universality of the representation space. Then TACO-DR introduces task sensitivity guided optimization (T-SAGE), which dynamically adjusts the gradient step size of each parameter according to its sensitivity to different tasks, to encourage parameters to capture more task specific signals.
To demonstrate the advantages of TACO-DR, we conduct experiments the eight retrieval tasks included in the standard KILT benchmark, including scenarios such as fact check, entity linking, slot filling, OpenQA, and dialog (Petroni et al. 2021). While standard multi-task dense retrieval failed to outperform their single task counter parts, TACO-DR outperforms both single task models as well as other more advanced multi-tass learning techniques. Our ablation confirms the source of effectiveness from its task identifier prompts and task-sensitivity guided optimization.
Our further analysis reveals how TACO-DR better leverages the model parameters for universal retrieval. Compared to standard multi-task learning, TACO-DR activates a larger fraction of the model parameters for each single task. effectively encourages more model parameters to capture task specific signals, which is achieved by enforcing the parameters to continuously focus tasks that they are initially learning during optimization. As a result, TACO-DR better utilizes the model parameters to capture the various training signals from multiple tasks and achieve strong retrieval accuracy with one universal retrieval system for many scenarios.
2 RELATED WORK
Learning text representations other than discrete bag-of-words. has been long desired goal in information retrieval (Deerwester et al. 1990 Huang et al., 2013). With the benefits of pretrained
language models (Kenton & Toutanova 2019]] Lee et al. 2019) and effective hard negative sampling
in finetuning (Karpukhin et al. (020aa] Xiong et al. 202 1, dense retrieval systems have shown strong
effectiveness on many retrieval scenarios (Qu et al. 2020 Herzig al., [021] Chang et al.. 2021).
Providing a centralized entry to multiple information sources has various advantages. It is more convenient for the user, reduces the information barrier between scenarios, and provides more diverse information access. Previously, this was achieved by complex divide-and-conquer systems, e.g., in federated search (Arguello et al. 2009). Universal retrieval provides a simpler solution with one unified representation space empowering multiple scenarios (Sciavolino) [02I), by embedding all data formats using neural encoders (Baevski et al. 2022).
Recent research universal retrieval often leverages multi-task learning to train one dual-encoder model for multiple retrieval tasks. Karpukhin et al. (2020a) train a DPR model on four OpenQA tasks and observed accuracy gains some tasks. Liu et al. (2022) shows the effectiveness of one unified embedding space to retrieve both texts and images. When the number of retrieval tasks grows, multi-task learning becomes challenging. Multi- trained DPR on eight KILT tasks (Petroni et al. I02I) does not provide much accuracy improvements over single task models (Maillard et al. 2021 Autoregressive retrieval, which directly generates document identifiers using query as input, is another potential solution (De Cao et al., 2020), but currently it mainly thrives on scenarios where the target corpus is small (Bevilacqua et al. 2022 or natural document identifiers exist (e.g., entity names (Chen et al. 2022)
Recent research has observed various optimization challenges in multi-task learning with modern deep learning systems (Ruder, 2017). for example, gradient conflicts, task imbalance, high variance of gradient magnitudes, to name few (Yu al.) (2200). Various techniques have been developed. A common way to address these challenges is to perform an operation on the gradients during scholastic updates, e.g., by projecting gradients from different tasks to their norms Yu et al., 2020), modifying their directions and magnitudes (Wang et al., 2020), and encouraging updates toward common directions (Piratla et al.).
3 PRELIMINARIES
In this section we recap preliminaries of dense retrieval and the constraints of universal retrieval.
Dense Retrieval. Given a query q and a corpus C, the retrieval task is to find a set of relevant documents d E C. for q, often by using a retrieval function f(g, d; 0). Dense retrieval (DR) systems
2

---

Under review as a conference paper at ICLR 2023
representationnccettic solution for retrieval (Lee al. 2019)
f(g, d; : sim 9 (d)
It first encodes d using the encoder g(: fqy) and o(: 0a), often finetuned from pretrained language models, e.g., BERT (Kenton Toutanova| 2019) and T5 (Ni et al. (022). The encoders can share the same parameter 0g = 02 use different weights. The similarity metric (sim() is simple, such as dot product and cosine which support efficient nearest neighbor search (Johnson et al. 2019).
The standard training in dense retrieval is to use NCE loss with sampled negatives Karpukkhin et al.
2020b Xiong et al. 2021):
exp(S d+; 0)) Σ log Ede JUEEGHILN exp(f(g, 4
T(())
The sampled negatives NEG dt) often consist of random negatives and hard negatives (top retrieved documents from retrieval model, for example, BM25 (Karpukhin et al. 2020b or dense retriever f() itself (Xiong al. 2021).)
Universal Retrieval aims to model multiple retrieval tasks T, each with different (but possibly overlapping) queries q¹ and corpus C", using one unified representation space (Maillard et 2021
The key of universality is to map queries and documents into the same embedding space:
Vi T f d²; sim(g' ; Og), g(d';'dd)))) 02)).
With different retrieval tasks sharing the same similarity metric in the unified space, universal retrieval provides unified information entry for users to conveniently access different resources, for example, one search place for web, email. enterprise documents, and multi media content. It also enables cross scenario retrieval, with relevance documents from the corpus of a different scenario than the query.
Another advantage of universal retrieval is efficiency. Building and maintaining a document index per scenario is costly, and sometimes infeasible with large scale corpora like the web. Unifying the docu- ment embedding space enables serving multiple scenarios with one index. To obtain this efficiency benefit, it is often necessary to use the same document encoder g(:;0a) 0a) across multiple retrieval tasks, as there is often overlap between task corpora. The query encoder is more flexible (Maillard ct al., 202I). Maintaining task- task-specific query encoders not as expensive and the query task is often available, e.g., designated by users predicted by query classifier (Li et al. 2008).
natural solution to universal dense retrieval is to train the encoders using multi-task learning:
exp(J d 0)) E d++) exp(f(g', d(''))
[(0) Σ w W'L'(0)= Σne log
ww
(4)
TET
TET
9²
combines the loss from all tasks and balances them by the hyperparameter w
ANALYSIS ON MULTI-TASK LEARNING IN UNIVERSAL RETRIEVAL
Ideally, the multi-tast -asti learning in Eqn. should provide effectiveness benefits with shared signals between tasks. However, recent research in universal retrieval often observed worse overall retrieval accuracy in comparison with task specific model [Maillard et al. 2021 Kojima et al., 2022). this section we conduct analysis to understand the challenges of multi-task learning in universal dense retrieval. We first briefly introduce the analysis setup, evaluate various multi-tate taski learning models, and then two methods to characterize the multi-task learned model.
Analysis Setup. We use the KILT benchmark, which includes eight tasks from five knowledge intensive retrieval scenarios and follow the exact settings from Maillard et al. (2021). The detailed information of KILT benchmark is listed in Appendix[A]
The results reported in previous research mainly finetune from BERT and its variants. To gather more observations, we train our own DR models using T5 (Raffel et al. 2019), which has shown better empirical results in search tasks (Ni al. 2022
Kojima et al. 2022
3

---

Under review as a conference paper at ICLR 2023
le7
Model
Avg R-Precision
NQ WoW
Reported in Maillard et al. 2021]:
Density
T-REx FEV
Density
DPR
67.19 / 42.48
MT-DPR
64.04 / 36.86
Our Models:
Per Task T5-ANCE
71.38 / 47.89
0.2
0.4 0.6
0.8
1.0
0.5
1.0 1.5 Entropy
2.0
QSpec-TS TS-ANCE
70.87 / 48.07
lc-7
MT-TS-ANCE
70.61 / 46.55
(a) Task Sensitivity
(b) Task Entropy
Table I: Average page- and passage-level
Parameter sensitivity analysis of MT-T5- ANCE. The task-specific sensitivity of the other four KILT tasks are in Figure 4 We drop outliers in (a).
Figure I:
R- R-Prrcision of single task retrievers versus multi-task universal retrievers on KILT.
Specifically, we use ANCE (Xiong et al., (021) to finetune the ST5-EncDec architecture (Ni et al. 2022). which takes the text sequences as input to T5's encoder and produces the sequence embedding using the first decoder output the special token: g(æ; 0) =T) We experiment with several different parameter sharing variants: Per Task T5-ANCE, the single task model, MT-T5-ANCE the multi-task model with only encoder for all, and OSpec-T5-ANC which uses task specific query encoders and a unified document encoder.
Multi-Task Performance. The performance of these variations listed in Table T All multi-task learned systems underperform single task models. Using a stronger pretrained model, T5, improves performance of all methods but the discrepancy of worse multi-task performance remains. Enriching the capacity of query encoder to task-specific also does not lead to stably better performance than Per Task T5-ANCE, The results indicate that model capacity may not be the bottleneck of universal retrieval. Increasing model capacity, either via a stronger language model or allowing task-specific query encoders, does not lead better performance than single task models.
Utilization of Network Parameters. Next we dive into the learned parameters of our multi-task model and check whether the multi-task training has effectively utilized existing model capacity. For this analysis, we adapt the sensitivity metric, which describes the importance of parameters a network (Molchanov et al., 2016 2019), to the multi-task setting. Specifically, we define two sensitivity-based metrics, task-specific sensitivity and task entropy, to characterize the learned parameters.
Task Task-Specific Sensitivity measures the importance of parameter to a specific task. The definition of sensitivity is the impact of a parameter w.r.t. a training target based on the changes of loss when the parameter is zeroed out (Molchanov et al. 2019).
Following recent research (Liang et al. calculate task-specific sensitivity:
2022), we use the first order Taylor expansion to efficiently
(T)(A)
I (6,)
00, x 0,).
(5)
calculates the sensitivity of 0, to the loss of task t. A higher T (0,) means the parameter has captured more signals for the task, as removing it leads to bigger losses on the task. A low sensitivity indicates the parameter is not utilized for the task.
Task Entropy describes the spread of parameter's sensitivity to different tasks:
exp(Ix(0.)/T)
5(0,) = - (r,),) log(p.(0)))
(c(..) = Ever cxp(Ie
(6)
first converts the task-specific sensitivity to the multinomial distribution (0.) using softmax and then calculates its entropy. The softmax temperature controls the sharpness of this distribution which we set as 2 to be consistent all our experiments. For active parameters (with non-zero sensitivity any task), a high task entropy, i.e., flat distribution pt (0.), indicates the parameter captures signals shared by multiple tasks. A low entropy reflects a sparse distribution, that the parameter carries signals dedicated to that task (Task-Specific).
4

---

Under review as a conference paper at ICLR 2023
Algorithm Task Sensitivity Guided Adaptive Learning (T-SAGE).
Require:
Flatten model parameter TR: ;; minibatches where each batch B E B further
divided by tasks B (Beher: moving average rate 3 E [0, 1]; temperature > 0; learning rate 77 E
Notations: median M - RT is the column-wis median; softmax R M
x the row-wise softmax; lrr is a vector [T] ones; O is the Hadamard product.
1: Initialize I + 0 = MM
2: for each batch = (Br in B do
3: Compute the task-specific loss (0) on Be for each E T.
4: Compute the gradient matrix G E R-xx/T (T) with each column G. (0)
5: Compute the sensitivity matrix I' Mx with each column 1; + G. O0. Normalize the sensitivity scales across tasks I' <- I' l'diag (median(("))
6:
7: Update the moving average I <- 31 + B)1'.
8: Update the parameter 0 t- 0 ( @ 1 where U sootmax(1/// ERMITTT
9:
end for
We leverage these tools to study the learned parameter of MT-ANCE-T5, multi-task learned model. The kernel density estimated distribution the learned parameters are in Figure [ From the plot make the following two observations of the multi-task learned model:
1. The Bad. The training left significant amount of model capacity the table. Figure la shows long tail distribution of task kesk-specific sensitivity, with majority of parameters under utilized for each task.
2. The Good. As desired, the multi-task learning pushes some parameters capture task spe- cific signals and some to capture shared signals. Figure reveals the bimodal distribution with two peaks towards each end of the entropy.
These observations indicate that major bottleneck in universal retrieval is that the optimization process does not fully utilize the network parameters, Our analysis also shows an encouraging sign that the parameters have the tendency to capture task-specific signals in multi-task learning, only not as effective as single-task models.
5 TASK- SAEEIAPYY OPTIMIZATION
In this section, we propose TACO-DR, which aims to unleash the model capacity for universal retrieval using TAsk speCific Optimization. TACO-DR first uses prompt convey the task identifier in the query. Then, motivated by our analysis in the last section, it employs task sensitivity guided adaptive learning (T-SAGE) that dynamically adjusts learning rates to encourage task specialty on parameters in multi-task learning.
Task Identification Prompts. TACO-DR uses the task name prompts (Raffel et al. 2019 in front of the query mark the task for the universal retrieval model:
g' Task Name [SP] q.
It concatenates (o) the task name string to the search query with a special token ([SP]). This is standard way in prompt-bass based learning to notify the model of the origin of the query. As we discussed in Section 3 the query side in universal retrieval is more flexible and can have different query formats for different tasks. The task prompts can customized by the user ("prompt engineering") or being predicted by query classifiers if needed. All the remaining parts of the model are kept the same as MT-TS-ANCE: T5 encoder shared for both query and document, dot product similarity, standard document inputs, and the same encoder model for all tasks.
Task Sensitivity Guided Adaptive Learning (T-SAGE). To encourage model parameters to capture task specific signals, TACO-DR uses the task sensitivity distribution to dynamically change the learning rate of each parameter during optimization. Recent research has shown the effectiveness of sensitivity guided adaptive learning in improving network parameter utilization (Liang al.
5

---

Under review as a conference paper at ICLR 2023
2022). T-SAGE extends it to multi-task learning by assigning adaptive learning rate for each parameter-task pair.
Formally, in the stochastic learning process, T-SAGE performs the k-th gradient step on 0 as:
(L*()*)
0 = 0² 21
Σ tET
(0.)
(8)
00.
where v) is the overall gradient step size, (0*) is the gradient of task parameter 0g, and u² (0;) is task specific learning rate for each parameter
As shown in Figure Ib] the multi-task learning. network parameters have the tendency to capture task specific signals shared signals. Ideally, T-SAGE uses u to encourage this tendency, for parameters who have a preference towards one task to commit on learning task specialty, and the rest to keep learning shared signals.
For that goal, we first calculate the task-specific sensitivity for parameters:
1 (0+) aL'(0*) (0*)
0.
(9)
O0k
7/5 (0.) - (0.)/Median) sm sm (1 (0m));
(10)
(0.) = .II (0,) + (1 -,1,1-1111
(11)
The per-step fask-specite sensitivity I, (et)) has high variance in stochastic training. Eqn. 10 normalizes sensitivity of parameter using the median of all parameters w.r.t. to this task. It mitigates the variances between different tasks. We use median instead of mean to account for the long tail distribution of task-specific sensitivity. Eqn. [I estimation. avoids sudden changes of task sensitivity based on one mini-batch.
introduces momentum mechanism to improve the
Then we use the parameter's task specialty distribution to define its dynamic learning rate:
(p(((T))))) w (0.) = Ever exp(1)
(12)
Eqn. dynamically sets the learning rate of parameter 0, for task i by comparing its task sensitivity it (0.) to other tasks. 0; more sensitive to a task I, it receives larger rate uf (0,) to capture more task specific signals. If the parameter has flat distribution, it receives balanced learning rate over multiple tasks, with the potential to learn shared signals between multiple tasks.
The emphasis task specialty is controlled by the softmax temperature т. smaller produces sharper distribution and emphasizes more on learning task specialty. We can first employ a "burn-in" period that trains with uniform learning rates or high temperature т, for parameters to declare their tendency. Then we can reduce the temperature to encourage the learning of task specialty and to better utilize the model parameters. The matrix version of T-SAGE is summarized in Algorithm for implementation references.
EXPERIMENTAL SETUP
We briefly describe experimental settings in this section. More details can be found in Appendix C
Datasets. We follow Maillard et al. (2021) and use eight tasks from KILT (Petroni et al. 2021). We randomly down sample the training data of the two largest datasets (T-REx and zskE) to the same order of magnitude as the rest. All the datasets share the same document corpus of 36 million disjoint 100-token Wikipedia passages preprocessed by Maillard et al. (2021).
Evaluation. We report page-level R-precision and passage-level R-precision. We use official KILT evaluation scripts to evaluate page-level R- precision and TREC Eva for passage level.
We derive in standard SGD. Changing to momentum hased method is straightforward. httpss:/ttret ni wtt..tto.tttteeeeeeeaeeeeee

---

Under review as a conference paper at ICLR 2023
Faet Check.
1.
Slot Filling
Open Domnain QA
Model
AV2
T-REx
2RR
NQ
TQA
WoW
Avg
BM25
50.13/ 40.06
3,47
58.60/
51.64
/5.98/
25,83/ 14.20
43.95/ 38.38
16.44/16.16 16.16
27.50/ 18.41
38.17
Autoregresivve Models with Large-Sized Models (For Reference).
MT- BART"
81,92/-
75.18/-
91.08/-
67.61
50.98
70
MT. Corpussrain
82.06/ /-
90.84
77.62/
98.26
59.10/-
50.07 1-
68.78 7+
53.75 /-
72.56/-
Dual Emcoder Dense Retrieval Models with Base-sized Models.
Task
73.60/43. /3.92
58.54
97.74/ 78.81
46.63/43.48
65.12
40.32/ 20.73
67.19/42..4
MT-DPP V2"
74,72/46696
K3.78
77.23/411 63,18/55,54
/41.7
61.51/ 28.80
44.21 38.42
61 61.95/224
39.70/ 24.07
64.04 36.86
Per Task ANCE
74,28//41.89
85.28
22,18/ 22.09
99.38/ 84.47
65.39/33...
16.79
69.08 22..57
53.63 / 27.61
71.38/
/47.89
QSpa--TS-ANCE
89.71/57.91
N7.56
72.72
60
85.15/71.37
64:01/ 32,92
49.74/45.50
29
55,03/ /2.
/48.07
MI-TS-ANCE
84.03/
95.62
70.96 66.
$6,04 / 69.79
62.46/ 30.51
49.78 44.35
06.04 0.004/25.0
19,95 / 32.55
70.61 /46.55
PCG-TS-ANCE
84.65 58.60
85.51
70.48 65.01
65.02/70.55
62,14/33..22
50.93/ 46.25
67.57
59.82/ 32.68
70.77 /47.48
CGD-TS-ANCE
83,15/54.95
79.03
62.32/
/43.26
65.18117.477
58.87/ 33.881
(43.91
GN-TS-ANC
84.10/57.399
85,18
70.08/ 64.45
81.94 68.61
62.71
29.43
49.09
1:009/433 /4.24
/5.03/
60.05 / 32.68
70.07 /45,84
TACO-DR
86.17/ 60.76 84.64
78.12/ 72.57
97.91 / 92.80
61,86/ 31.16
50,61 46.72
09.62/28.32
60.97/33.24
73.74/ 55.99
Table 2: Page/passage level R-precision on KILT validation data. Bold indicates the best dual encoder model and underline indicates the second. Only page-level retrieval is defined for AIDA-YAGO (AY2). i and mark results from Chen et al. (2022) and Maillard al. (222I.
Model configuration. We use the standard dual-encoder architecture, initialized with T5-base (Raffel et al. 2019) and take the first decoder output as the sentence embedding (Ni et al.] (022). The query encoder and passage encoder share weights in all our runs except QSpec-TS-ANCE.
Training details. We train MT-T5-ANCE from T5-base using vanilla multi-task learning, for 20 epochs using BM25 negatives (Karpukhin et al., 2020b) and then seven ANCE (Xiong et al. 2021) episodes with its own hard negatives refreshed at the beginning of each episode. We use similar process for TACO-DR but with uniform weights (infinite softmax temperature) in T-SAGE for the first six ANCE episodes and 2 in the last episode. Training TACO-DR with annealed temperature earlier in the process achieves similar results. We provide more details in Appendix [ We use Adam (Kingma & Ba 2013) with linear learning rate decay schedule with warmup proportion 0.1 over epochs for each ANCE iteration.
Baselines. We compare our model with recent universal retrieval models (Maillard et al., 2021 Chen et al. 2022). We also compare with general multi-task learning models PCG (Yu et al. (020). CGD (Piratla et al. 2022) and GN (Chen et al.) 2018) (using our own implementations). A parallel work uses autoregressive models to generate entity names (as document identifier) to conduct retrieval on KILT CChn et al. 2022). They use large size models and large amount of additional training data. We list their models trained solely on KILT training labels for reference.
7 EVALUATION RESULTS
This section presents four experiments on TACO-DR's retrieval accuracy, its ablations, learned task specialty, and its parameter commitment on tasks.
7.1 OVERALL RETRIEVAL ACCURACY
Table lists the retrieval accuracy of evaluated methods. We evaluate R-precision at both page and passage level but will mainly focus on page level evaluation later experiments, which is the metric designed for many tasks (Kwiatkowski et al., 2019 Nie et al. 2019) e.g.).
TACO-DR outperforms all other methods on the average accuracy. It achieves best single task page- level R-precision on four of the eight KILT tasks, the most robust among all evaluated methods. TACO- DR is the only evaluated universal retrieval method that outperforms its task-specific counterpart, Per Task T5-ANCE, on both page and passage level average accuracy. Though the autoregressive models employ a large sized pretrained model and benefit from the existence of semantic document identifiers (entity name) in KILT, TACO-DR achieves better overall accuracy with smaller model when all are trained on KILT data.
Recent multi-tast task learning techniques, PCG et. 2020), CGD (Piratla et al. 2022) and GN (Chen ct al. 2018), do not outperform standard multi-task learning in universal retrieval. These methods mainly focus on mitigating the conflicts of gradients from different tasks. Our analysis shows the
7

---

Under review as a conference paper at ICLR 2023
Faet Check.
Ent. L.
Slot Filling
Open Domain QA
Dial.
Variants
FEV
AY2
T-REx
zsRE
NQ
HoPo
TQA
WoW
Avg
MF-TS-ANCE
84,03
85.62
70.96
86.04
62.46
49.78
66.04
59.95
70.61
QSpec-TS-ANCE
82.71
87.56
72.72
85.15
64.01
49.74
69.12
55.93
70.87
MT-TS-ANCE Marker
84.49
85.51
73.88
89.37
62.85
50.97
67.70
60.02
71.85
TACO-DR w.o. T-SAGE
84.81
85.49
75.00
92.24
62.81
51.47
68.95
60.54
72.66
TACO-DR w.o. Prompt
85.71
74.82
94.68
61.05
49.38
67.79
58.81
72.12
TACO-DR
86.17
84.64
78.12
97.91
61.86
50.61
60.97
73.74
Table 3: Page level accuracy of TACO-DR variants. MT-T5-ANCE, w. Marker uses the coarse task marker from Maillard al. (202T). Best results are bolded.
lc7
Je7
1.777
le7
ME-IS-ANCE:
MI-TS-ANCE
ME-TE-ANCE
TACE-DR wa. 1-SAGE
TACO-OR vu T-SAGE
TACO-DR wa. T-SAGE
TACCSIR WAs T-SAGE
TACO-D8. Prongi
TACO-UR
TACO Pounpi
JACC-UK Prorpt
Density
Deneity
Density
Density
0.5
0
1.5 3.0
4.5
6.0
1.5
3.0
4.5
6.0
1.5
3.0
4.5
6.0
1.5
3.0
4.5
6.0
1o-8
Ic-8
lc-8
1c-8
(a) NQ
(b) WoW
(c) T-REx
FEV
Figure 2: Task-specific sensitivity distribution on the training data of four KILT tasks. The final models are used. The distribution of the other four tasks are shown in Figure O in Appendix. The x-axis is sensitivity, and we drop outliers that are far from the median to case visualization.
challenge is mainly on under captured task specific signals and TACO-DR is designed to increase the utilization and task specialty of model parameters. Section 7.3 further analyzes this advantage of TACO-DR.
Note that no method performs the best on every retrieval task. Universal retrieval is challenging with scenarios consisting of variable training data quantity, task difficulty, and different user preferences. How to achieve robust effectiveness benefits a next step of development for universal retrieval.
7.2 ABLATION STUDIES
The page level results of TACO-DR variants and corresponding baselines shown Table
Using Task Identification Prompt in standard multi-task learning also improves its performance. TACO-DR w.o. T-SAGE outperforms MT-T5-ANCE w. Marker where the only difference is the task prompts. Both outperform QSpec-TS--S T-ANCE which uses significantly more parameters, further confirming that the model capacity is the biggest bottleneck in universal retrieval.
Both task prompt and T-SAGE contribute to the effectiveness of TACO-DR. The two target different parts of multi-task learning. The task prompt conveys the task specific information in model inputs and T-SAGE adjusts the learning process to emphasize task specialty during optimization. The benefits of prompt are more well understood in multi-task learning (Raffel et al. (019)) and in the next experiments we focus on the properties of T-SAGE.
7.3 PARAMETER UTILIZATION AND TASK SPECIALTY
In this experiment we analyze TACO-DR's parameters using task-specific sensitivity and task entropy.
Figure plots the kernel density estimated distribution of task-specific sensitivity in TACO-DR and its variants. There is a significant reduction of non-activated parameters in TACO-DR compared with MT-T5-ANCE This mainly attributes T-SAGE. Using with or without prompt, significantly reduces the peak on the low sensitivity side, showing better utilization of parameter capacity for each task (Liang al. 2022).
8

---

Under review as a conference paper at ICLR 2023
ME-TS-ANCE
80.0
30
TACO-DR wAs T-SAGE
AVG R-Precision
AVG R-Precision Special Fraction
Special Fraction
Percentage
TACO DR wo. Pronpt
77.5
0.22
20
75.0
0.21
10
72.5
Task
10.3,
10.66, 11388 11.74,
70.0
0.1
2.0
10
100
0.20
Special0.66)
1.02) 1.38) 1.74) 2.1
Entropy
(a) Task Entropy.
(b) Influence of T-SAGE Temperature.
Figure 3:
The analysis of task specialty on parameters learned by different methods (a) and when using different temperatures in T-SAGE.
Figure 3a plots the histograms of task entropy for the learned parameters. We first group parameters into two special bins. The first "Task Specific" bin that includes parameters whose entropy is smaller than 0.3, which is the entropy of 95% probability on one task and the 5% uniformly on the rest seven. The "Not Activated" bin includes parameters whose sensitivity w.r.t tasks is near zero (< 0.1). TACO-DR significantly improves the fraction of task specific parameters to 22%, in comparison with 19% in MT-TS-ANCE. also reduces the fraction of not activated parameters, showing optimizing task specialty also better utilizes the model capacity.
Figure 3b illustrates the influence of the softmax temperature T in Eqn. 12 which controls the sharpness of the sensitivity distribution thus the encouragement of task specialty. The accuracy of TACO-DR is quite stable w.r.t. different T values in normal temperature range. At the same time, does show the effectiveness of in controlling for different levels of task specialty. A smaller leads to more task specific parameters and a bigger 7 leads to fewer ones.
7.4 COMMITMENT OF TASK SPECIFIC PARAMETERS DURING OPTIMIZATION
This last experiment studies the mo- mentum mechanism in T-SAGE. Be- sides retrieval accuracy, we also calcu-
0
0.6
0.7
0.8
0.9
0.999
Avg R-prec
72.91
72.98
73.02
73.16
73.61
73.74
late the task commitment rate, which
Step
Task C C'ommitment Rate
fraction of task-specific parame-
400
43.8%
41.2%
41.4%
39.7%
36.5%
35.5%
800
61.6%
57.8%
58.8%
56.4%
56.3%
52.0%
ters that are committed to the same
1200
63.4%
58.5%
60.0%
59.5%
61.1%
56.4%
task after some training steps (400),
68.7%
65.4%
68.3%
70.4%
70.2%
ie., being most sensitive to the same task. The results in Table
2000
72.5%
74.3%
73.8%
72.5%
75.9%
76.3%
Table 4: Impact of momentum decay factor 3. The commit- ment rate calculated by comparing the task distribution at the corresponding step to the previous row.
At the early stage of optimization, the task commitment rates are low. Pa- rameters are still wandering around signals from different tasks. Strong momentum factor carries more random signals from early chaotic stages and in fact hurts the com- mitment rate. With enough training. parameters start to settle to their target tasks and we observe significant increments of commitment rates, even for the ones without momentum 3 = 0. There stronger momentum provides better commitments and eventually leads to better retrieval accuracy.
8 CONCLUSION
In this paper we present TACO-DR that unleashes the model capacity for universal retrieval by encouraging parameter task specialties in multi-task training. We first analyze the standard multi-task learning in universal retrieval and show that much of the model capacity is under utilized. TACO-DR dynamically adjusts the learning rate of each parameter for each task based on its task sensitivity, to encourage the learning of task specific signals during stochastic training. Our experiments demonstrate the benefits of TACO-DR, its optimization task specialty, better utilization of network parameters, and effective control of the optimization process.

---

Under review as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
To enhance reproducibility, we describe our model configuration and overall training details in Section Beyond of that, we include more details and hyperparameters in Appendix. Appendix [A] shows data statistics and data-related hyperparameters. Appendix proovdee batch sampling strategy, training hyperparameters and more training details. We plan to submit our code after the discussion forums open. We will make a comment containing a link of our anonymous repository internally visible to reviewers and ACs. We will release our code and model checkpoints, along with analysis scripts if this work is accepted.
REFERENCES
Jaime Arguello, Fernando Diaz, Jamie Callan, and Jean-Francois Crespo. Sources of evidence for vertical selection. In Proceedings of the 32nd international ACM SIGIR conference on Research and development information retrieval, pp. 315-322, 2009.
Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli. Data2vec: A general framework for self-supervised learning in speech, vision and language. arXiv preprint arXiv:2202.03555. 2022.
Payal Bajaj. Daniel Campos, Nick Craswell, Li Deng. Jianfeng Gao, Xiaodong Liu, Rangan Ma- jumder, Andrew MeNamara, Bhaskar Mitra, Tri Nguyen, et al. Ms marco: A human generated machine reading comprehension dataset. In Proceedings of CoCo@ NIPS 2016, 2016. URL 2016 paper9. pdf
http::/ceuu-- ws. http:///eeeeassssoo///////7/1110
Michele Bevilacqua, Giuseppe Ottaviano, Patrick Lewis, Wen tau Yih, Sebastian Riedel, and Fabio Petroni. Autoregressive search engines: Generating substrings as document identifiers. In arXiv pre-print 2204. 10628, 2022.
Yinghsan Chang, Mridu Narang, Hisami Suzuki, Guihong Cao, Jianfeng Gao, and Yonatan Bisk. WebQA: Multihop and Multimodal QA. 2021. URL
00590
Jianuui Chen, Ruqing Zhang, Jiafeng Guo, Yiqun Liu, Yixing Fan, and Xueqi Cheng. Corpusbrain: Pre-train a generative retrieval model for knowledge-intene intensive language tasks. arXiv preprint arXiv:2208. 076552 2022.
Zhao Chen, Vijay Badrinarayanan, Chen- Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In ICML, 2018.
Nicola De Cao, Gautier Izacard. Sebastian Riedel, and Fabio Petroni. Autoregressive entity retrieval. arXiv preprini arXiv:2010.00904 2020.
Scott Deerwester, Susan T Dumais, George W Furnas, Thomas K Landauer, and Richard Harshman. Indexing by latent semantic analysis. Journal of the American society for information science, 41(6):391 407, 1990. URL http:tttp///iteeeeeeeeee 0oi=10.1.1 52. 112.1111ree-reeeltyyppaadd
ist. psu. edu http.eeeee.leeedo//ooo..ooo
Jonathan Herzig, Thomas Müller, Syrine Krichene, and Julian Martin Eisenschlos. Open domain
question answering over tables via dense retrieval. arXiv preprint arXiv:2103. 120111, 2021.
Po-Sen Huang, Xiaodong He, Jianfeng Gao, Li Deng, Alex Acero, and Larry Heck. Learning deep structured semantic models for web search using clickthrough data. In Proceedings of CIKM 2013, pp. 2333-2338, 2013.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scall scale similarity search with gpus. IEEE Transactions on Big Data, 5(3:5335-547, 2019.
Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 2769-6781 Online, November 2020a. Association for Computational Linguistics. doi: 10. 18.1386//vlll 102.000011/020.mmll.mm..mm m50n.550. URL https: www. ac wweb. org/ httgg//nthoooy.../
2020 emn 1p-ma in
10

---

Under review as a conference paper at ICLR 2023
Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNL.P, pp. 6769-6781, 2020b.
Jacob Devlin Ming-Wa Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HILT 2019, pp. 4171-4186, 2019. URL
14223
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. arXiv preprint arXiv:2205. /916, 2022.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research, Transactions of the Association for Computational Linguistics, 7:453-466, 2019.
Kenton Lee, Ming-Wei Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. In Proceedings of ACL 2019, pp. 6086-6096, 2019. URL
bttps::/aclaatthooorry. ott.:/oooy...o1--1112201
Patrick Lewis, Ethan Perez, Aleksandara Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmented genera- tion for knowledge- knowledge-intensiee intengivee nlp tasks. arXiv preprint arXiv:2005. 1401. 2020.
Xiao Li, Ye-Yi Wang, and Alex Acero. Learning query intent from regularized click graphs. In Proceedings of the 31 annual international ACM SIGIR conference on Research and development in information retrieval, pp.
2008
Chen Liang, Haoming Jiang, Simiao Zuo, Pengcheng He, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Tho Zhao. No parameters left behind: Sensitivity guided adaptive learning rate for training large transformer models. In International Conference on Learning Representations, 2022. URL /tttps://cpenreeieee net. wttp://wwwwweeeeeeeeeneeeeed www.oooodocuuueeoooo. __VVN
Zhenghao I Liu, Kaitao Zhang, Chenyan Xiong. Zhiyuan Liu, and Maosong Sun. Openmatch: An open source library for neu-ir research. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 2531-2535, 2021.
Zhenghao Liu, Chenyan Xiong, Yuanhuiyi Lv, Zhiyuan Liu, and Ge Yu. Universal multi multi-modality retrieval with one unified embedding space. arXiv preprint arXiv:2209.001 2022.
Jean Maillard, Vladimir Karpukhin, Fabio Petroni, Wen-tau Yih, Barlas Oguz, Veselin Stoyanov, and Gargi Ghosh. Multi-task retrieval for knowledge-intensive tasks. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing ( (Volume 1: Long Papers), pp. (098-1111) 2021.
Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convolutional neural networks for resource efficient inference. arXiv preprint arXiv:1611.0664 2016.
Pavlo Molchanov, Arun Mallya, Stephen Tyree, Turi Frosio, and Jan Kautz. Importance estimation for neural network pruning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1164-11222, 2019.
Jianmo Ni. Gustavo Hernandez Abrego, Noah Constant, Ji Ma, Keith Hall, Daniel Cer, and Yinfei Yang. Sentence-t5: Scalable sentence encoders from pre- tre-trained text-to-text models. In Findings of the Association for Computational Linguistics: ACL 2022, pp. 1864-1874, 2022.
Yixin Nie, Songhe Wang, and Mohit Bansal. Revealing the Importance of Semantic Retrieval for Machine Reading at Scale. In Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, 2019.
11

---

Under review as a conference paper at ICLR 2023
Fabio Petroni, Aleksandra Piktus, Angela Fan, Patrick SH Lewis, Majid Yazdani, Nicola De Cao, James Thore. Yacine Jernite, Vladimir Karpukhin, Jean Maillard, et al. Kilt: benchmark for knowledge intensive language tasks. In NAACL-HLT, 2021.
Vihari Piratla, Praneeth Netrapalli, and Sunita Sarawagi. Focus on the common good: Group distributional robustness follows.
Vihari Piratla, Praneeth Netrapalli, and Sunita Sarawagi. Focus the common good: Group distributional robustness follows. In International Conference on Learning Representations, 2022.
URL httphhtts///pppeeee
wwt.:ewweeee...eeeeeeeee id=i FARV 2
Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Wayne Xin Zhao, Daxiang Dong, Hua Wu, and Haifeng Wang. Rocketqa: An optimized training approach to dense passage retrieval for open-domain question answering. arXiv preprint arXiv:2010.08191, 2020.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 2019.
Sebastian Ruder. An overview of multi-task learning in deep neural networks. arXiv preprint arXiv:1706.05098, 2017.
Christopher Sciavolino. Towards universal dense retrieval for open-domain question answering. arXiv preprini arXiv:2109./ 1085, 2021.
Zirui Wang, Yulia Tsvetkov, Orhan Firat, and Yuan Cao. Gradient vaccine: Investigating and improv- ing multi-task optimization in massively multilingual models. arXiv preprint ar.10:2010.05874 2020.
Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul N. Bennett, Junaid Ahmed, and Amold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. In International Conference on Learning Representations, 2021. URL https //openreview net. wtt.:e/eooommmmeeeerrrrrrr
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 33: 5824-5836, 2020.
A DATA STATISTICS
See Table for data statistics and some data-related hyperparameters. We randomly downsample T-REx and zsRE to bring them to the same order of magnitude as the others.
Dataset
# #Train
Batch size
Max query length
Natural Questions
76445
16
32
TriviaQA
52886
14
HotpotQA
68659
15
36
Wizard of Wikipedia
79535
16
256
T-REx
94514
16
32
FEVER
70757
15
Zero Shot RE
99500
17
32
LIDA-YAGE 2
17895
128
Table 5: Data statistics and some data-related hyperparameters for our experiments. The original T-REx and zsRE have very large dataset size. We randomly downsample T-REx and zsRE to bring them to the same order of magnitude as the others.
12

---

Under review as a conference paper at ICLR 2023
Fact Check.
Eat. L.
Slot Filling
Open Domain QA
Dial.
Model
FEV
AY2
T-REx
ZRRE
NQ
HePo
TQA
WoW
Aug
Annealed TACO-DR
/9.99
84.26
72.155
97.15/8228
61.30/ 30.45
50.60/ 46.69
69.34/ 27.91
61.43/33.6
73:47/ 50.45
TACO-DR
66.17/60.76
84 64
78
78.12/72.51
97.91 18280
61.86/ 16
50.61 /46.72
69.62 26.32
60.97/33.88
73.74/50.00
Table 6: Page- and passage- level R-precision for TACO-DR with annealed temperature started from 4-th ANCE episode with annealing schedule т 64 (0.5)a where >= is the ANCE episode number. We also include the performance of TACO-DR trained with the setting in section 6 reference.
for
learning rate
warmup ratio negatives per 0.1
epochs T
total batch size
5e-6
2.0
0.999
120
Table 7: Training hyperparameters for training our TACO-DR model. We use Ainmmingma & Ba (2015) with learning rate 5e 6. We use linear learning rate schedule with warmup raio 0.1. Each query uses hard negatives for training. Each ANCE episode trains for 3 epochs. Total batch size of all task batches are 120.
B OTHER PLOTS
Figure 4 shows parameters sensitivity distribution of the remaining tasks TQA, HoPo, zsRE and AY2 that Figure 2 doesn't show. We observe similar trend to Figure 2 that T-SAGE significantly reduce non-activated parameters except for zsRE. We note that zsRE a very simple task that can achieve almost 100% accuracy, which doesn't 't require many activated parameters.
1.5 1o7
2.0
1e8
107
ME-15-ANCE
MIRTS-ANCH
MI-TS-NNCE
TAUU-DI wa. ESAGE
TACISIR vO. T-MALE
TACO DR T-SAGE
LACO-IIC was 1-AALE
Donsity
TACO-DR
Density
TACO-OR
Density
TACO v.o. Pocmut
Donsity
1.0
TACO-DR was. Prinqe
1.0
0.5
0.5
1.5
4.5
6.0
3.0
4.5
6.0
0.4 0.6 D.
1.5
3.0
4.5
le 8
1c 8
le
1e
(a) TQA
(b) HoPo
(c) zsRE
(d) AY2
Figure 4: Parameters sensitivity distribution using kernel density estimation on the training data of the remaining four KILT tasks that Figure doesn't 't show.
C MORE EXPERIMENT DETAILS
Batch Sampling We follow Raffel et al. 2019) and temperature-scaled mixing sampling strategy compute batch size for each task t T: B, x (N:/ Eeer N/.) for some temperature e (we set it to 4 in our experiments). Here N₂ is the dataset size of task t. Note that we compute task loss of each task batch independently instead of mixing all task batches for every optimization step. Each dataset needs to sample different number of batches to cover every training sample in that dataset once. We the maximum of them the number of batches that every dataset needs sample. We shuffle and cycle batch sampling iterators of datasets that finish iterating early. Batch size of each dataset computed by setting mixing temperature 4 and Σ Ever N₂ 120 is in Table
Other training details The data-related hyperparameters, such maximum input query length and batch size, are listed in Table S The training hyperparameters are listed in Table 7 We use NCE loss with cross device in-batch negative mixed with hard negatives to compute each task loss. We sample two hard negatives for each query. We employ a "burn in" period for the first 10% training steps with uniform learning rates for parameters declare their tendency during T-SAGE optimization. All of our experiments are run on a machine with 8 A100-80GB GPUS. Our implementations are built upon Open (Liu et al. 2021).
13

---

Under review as a conference paper at ICLR 2023
Annealed Softmax Temperature We also try annealing the softmax temperature in T-SAGE early at the beginning of the fourth ANCE episode. Specifically, we anneal the softmax temperature according to schedule T, = 64 (0.5)*-4 where >= 4 is the a-th ANCE episode. The results listed Table We achieve similar performance as the setting in section 6 The training setting in section 6 has low cost with slightly higher performance compared to this annealed setting.
14


Reviews:
[{"review_id": "Xf-b7P6vmAS", "reviewer": "Reviewer_ZTG9", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies the performance of multi-task learning in universal dense retrieval systems. The authors show that the network capacity is not the limiting factor for the universal retrieval accuracy. Instead due to insufficient optimization, large portion of parameters show low sensitivity to each task. A method, called TACO-DR, is proposed that optimizes the task-specialty of the parameters. ", "summary_of_the_review": "Overall, I found the idea of having task-specific adaptive learning rate interesting. However, I believe the empirical results can be further improved.", "strength_and_weaknesses": "The paper is generally well-written and easy to follow. I have the following comments and questions:\n- If I understand correctly, the results in table 2 of the paper should be comparable to table 3 of Chen et al. (2022). It seems that their best results are slightly different from what is reported in table 2. For instance, for T-REx and TQA Chen et al. (2022) are reporting 85.03 and 71.71 while the results in table 2 are 77.62 and 68.78. If possible and in order to make the results more comparable, can you use a similar training strategy to theirs? \n- For the sake of comparing the models on all 11 KILT tasks, can you also add the results for the 3 remaining tasks? \n- Can you provide the confidence intervals in Figure 3(a)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "Tn6Uem8Wsmw", "reviewer": "Reviewer_6W5P", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The submission concerns universal retrieval, or multi-task dense retrieval, where queries and documents from different sources are mapped to the same space. The paper first conducts a study using T5 on the KILT dataset and empirically shows that there is some gap between single task performance vs multi-task. Then parameter importance analysis is done using existing techniques, showing that it is mainly due to parameters under capturing task specific signals. Then an optimization with adaptive learning rates is applied. Experiments on the KILT shows the proposed framework can achieve better average performance than alternatives, including BM25, naive multi-task learning, and two multi-task methods can focus on mitigating gradient conflicts.", "summary_of_the_review": "Overall the paper deals with a relatively new problem and is clearly written. The novelty is a major concern for a top-tier ML venue (an NLP venue may make more sense). Please see the detailed summary and discussions above. \nSome other comments:\n\n“The results indicate that model capacity may not be the bottleneck of universal retrieval” - what size of T5 is used (this should be clarified anyway - looks like BASE is used from experiments)? It looks like the gap is smaller in Table 1. What if larger T5 models are used?\n", "strength_and_weaknesses": "Strength\nThe paper is generally well written, though some details can be provided more properly.\nThe studied problem is relatively new.\n\nWeakness\nNovelty. The analysis and method used in this paper are applications of existing methods. Though it is a good effort, this is unlikely to be sufficient for a top-tier ML venue.\nGiven the rich literature of multi-task learning, there are methods that do not focus on gradient conflicts. The chosen related baselines are likely to be used in the wrong context.  \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "EC5EuNxyEc", "reviewer": "Reviewer_KDMG", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work proposess a multi-task universal retrieval model that is based on task-specific prompt tuning and adaptive learning are based on task sensitivity. Experiments on the KILT benchmark dataset over existing models, and ablation on the task sensitivity, show that although the model capacity of existing works is not sufficiently utilized, universal retrieval models can capture task-specific signals. The proposed method tries to address the limitation of model capacity utilization.", "summary_of_the_review": "The analysis of existing models is promising and could lead to new universal retrieval methods. However, the current proposed method seems to be a direct application of existing methods in the literature. The introduction and related work could be enhanced to clearly showcase any technical contributions. The proposed work could be further improved by comparing the task-sensitivity adaptive learning rate method with other adaptive learning methods in the literature.", "strength_and_weaknesses": "[+] The analysis of task sensitivity provides a clearly articulated motivation for the proposed work.\n\n[-] In many cases, the performance difference between TACO-DR and the next best baseline is marginal.\n\n[-] The proposed work seems to be a straightforward application of prompt tuning and task sensitivity adaptive learning rate (similar to momentum-based optimizers in the literature). If there exist novel unique contributions, these could be better highlighted in the introduction and related work.\n\n[-] More relevant baselines are needed, e.g., other formulations of task-specific learning rates, e.g., a self-paced version that the task-specific hyper-parameter is weighted based on a loss-based task difficulty score, or an entropy-based task-specific learning rate. If baselines are trained with vanilla gradient descent, it would be interesting to see if standard Adagrad produces the same results. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "2e6wDsp9Pdl", "reviewer": "Reviewer_UpYV", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors introduce a multi-task learning method for universal retrieval. They identify that the reason existing approaches fail is because they fail to capture task-specific signals and therefore propose a method that learns task-specific parameter-specific learning rates to better capture task-specific signals. ", "summary_of_the_review": "The presented work does achieve state-of-the-art results for multitask retrieval. TACO-DR is an extension of an existing method (Liang et al., 2022). Although it works well for retrieval, it's not clear whether this method is useful just for this one setting or is more broadly useful. It would be interesting to investigate whether this method provides benefits for other multi-task learning applications. ", "strength_and_weaknesses": "Strengths:\n* The authors extend sensitivity guided adaptive learning (Liang et al., 2022) to the multi-task setting.\n* Empirically, the authors show that they are able to outperform multi-task and single-task baselines on the KILT benchmark.\n* The authors also show that adding task labels as a prompt improves multi-task retrieval performance. \n\nWeakness:\n* The sensitivity analysis graphs in  figure 1 and figure 2 are unclear and need more explanations. The x-axis in unlabelled in the graph and it is unclear what the axes are and what range of values they show. In figure 2, is there a tail that is not shown. How come the peak on the left is lower for the TACO methods but the tail is the same for all the methods?\n* The percentage ranges for the values in figure 3a and table 4 don't seem to be consistent with each other. Is the commitment rate calculated differently for these?\n* It's not clear whether the current framework can handle the case when a parameter is important for a subset of the tasks but not all the tasks. \n* This method doesn't seem specific to retrieval. Can other multitask settings also benefit from using task-specific parameter-specific learning rates?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject