Multi-Task Pre-Training for Plug-and-Play Task-Oriented Dialogue System

In plain words

Many task-oriented chatbots, such as a booking assistant, are built as a chain of steps. One step works out what the user wants, the next decides what the system should do, and the last writes the reply. Each step feeds the next, so a mistake early on carries forward. Building these chains also needs training data labelled for every step, which is expensive and rules out many existing conversation datasets that are labelled for only some steps.

The paper's idea is to use one language model for all the steps. Each step becomes a sentence-in, sentence-out task: the model gets a short instruction that names the step, such as "translate dialogue to belief state:", followed by the conversation so far, and writes the answer as text. The model is first pre-trained on many existing dialogue datasets, each labelled for only some of the steps, so it can use whatever labels a dataset has. Then it is fine-tuned on a target task with the same objective. The authors report that this beats earlier systems on three benchmark tasks and is especially strong when little labelled data is available.

TL;DR

• One T5-based text-to-text model covers intent classification, belief-state tracking, dialogue action prediction and response generation, selected by a task prompt §3.2 fig. 1.
• Pre-training uses eleven human-written multi-turn dialogue corpora with over 2.3M utterances, each annotated for only some tasks §3.1 tab. 1.
• Reported MultiWOZ 2.0 end-to-end: combined score 102.92 for PPTOD-base, the top value in the paper's table tab. 2. Author-reported.
• Reported low-resource DST (1% of MultiWOZ 2.0 training data): joint accuracy 31.46 for PPTOD-large against 13.21 for SOLOIST tab. 5.
• The released code is Apache-2.0. It matches the paper's batch size, learning rate and epoch count, but it differs in loss scaling, data shuffling and sequence truncation (see Pitfalls).

Problem

Most earlier neural task-oriented dialogue systems solve belief-state tracking (DST), action selection (POL) and response generation (NLG) in a fixed order, so each step conditions on the previous steps' outputs §1. The authors name three costs of this design: errors from early steps propagate §1; the training data must be annotated for every sub-task, which blocks partially annotated corpora §1; and the sequential order adds inference latency §1. Earlier dialogue pre-training methods need corpora annotated for every sub-task §2.

Method

The paper casts every task-oriented dialogue sub-task as text generation with a plugged-in prompt. It pre-trains a T5 model on mixed partially labelled data, then fine-tunes it with the same objective §3.2 §3.3 fig. 1.

1. Sample format. Each training sample is a triple of task prompt, dialogue context and target text, written d = (z_t, x, y) §3.2 eq. 1. The prompt z_t has the form "translate dialogue to A:", where A is "user intent", "belief state", "dialogue act" or "system response" for NLU, DST, POL and NLG §3.2. The context x is every earlier utterance, from both system and user, in order §3.2. The target y is the expected output text for that task, for example the intent label §3.2.
2. Pre-training corpora. Eleven corpora are combined: MetaLWOZ, SNIPS, CLINC, ATIS, KVRET, WOZ, CamRest676, MSR-E2E, Frames, TaskMaster and Schema-Guided §3.1 tab. 1 app. A. Each corpus is annotated for some of NLU, DST, POL and NLG, and a sample is produced only for the annotations a corpus has tab. 1. The authors report over 2.3M utterances §3.1.
3. Objective. The model is trained by maximum likelihood on the target tokens, conditioned on prompt and context §3.2 eq. 2. Training uses mini-batches over the mixed data alg. 1.
4. Fine-tuning. For a new task, the pre-trained PPTOD is trained with the same objective on that task's labelled data §3.3 eq. 2.
5. Inference, end-to-end. The model predicts the belief state first. The predicted state is used to query the database. The dialogue action and the response are then produced in parallel, both from the context and the database result §4.1.1.
6. Inference, DST and intent classification. The model generates the state or the intent label directly §4.2 §4.3.

# Build one training sample d = (z_t, x, y)  [§3.2] [eq. 1]
# Output: source text and target text

def build_sample(task, dialogue_turns, target_text):
    # Prompt names the task, e.g. "translate dialogue to belief state:"  [§3.2]
    prompt = TASK_PROMPTS[task]
    # Context: all earlier turns, system and user, in order  [§3.2]
    context = join(dialogue_turns)
    # Target is the expected output text for this task  [§3.2]
    return prompt + context, target_text

# Pre-training loop  [alg. 1]
# Input:  mixed dataset of samples; number of epochs
# Output: trained model parameters

def pretrain(model, dataset, epochs):
    # Repeat over the whole mixed dataset  [alg. 1]
    for epoch in range(epochs):
        # Shuffle; the paper says data from different tasks is mixed  [alg. 1]
        samples = shuffle(dataset)
        # One mini-batch at a time  [alg. 1]
        for batch in make_batches(samples, size=128):
            # Token-level maximum likelihood on the target  [eq. 2]
            loss = token_nll(model, batch)
            update_parameters(model, loss)
    return model

# End-to-end inference  [§4.1.1]
# Input:  dialogue context and database
# Output: system response

def respond(model, context, database):
    # Predict the belief state from the context  [§4.1.1]
    belief = generate(model, bs_prompt(context), max_len=120)
    # Query the database with the predicted state  [§4.1.1]
    db_result = database.query(belief)
    # Action and response both take the DB result as input  [§4.1.1]
    # [code: E2E_TOD/inference_utlis.py#L137]
    # [code: E2E_TOD/inference_utlis.py#L154]
    action = generate(model, da_prompt(context, db_result), max_len=80)
    response = generate(model, nlg_prompt(context, db_result), max_len=120)
    return response

The paper's Eq. 2 is a sum of token negative log-likelihoods eq. 2. Algorithm 1 writes the same loss as a mini-batch update alg. 1. Neither states the batch normalisation or the decoding caps. Both are taken from the code in Pitfalls.

Figure 1 shows the four pre-training tasks. Each receives a prompt and the dialogue context and generates target text. It also shows that training uses partially annotated data across the tasks fig. 1.

Hyperparameters

Pre-training values are from the paper §3.4 unless marked as code.

Results

Author-reported. The table gives the paper's MultiWOZ 2.0 and 2.1 end-to-end results: Inform, Success, BLEU and Combined score. Setup: full training data, one run per row §4.1.3 tab. 2. Rows marked † need oracle dialogue states at each turn; rows marked ‡ use author-released models tab. 2.

TOP and TOP+NOD results are cited from the original paper, not re-run tab. 2. Source: tab. 2. The paper notes that PPTOD-large is weaker than small and base on this task, and attributes it to delexicalised tokens that are unseen in pre-training §4.1.3.

Low-resource end-to-end (MultiWOZ 2.0), Combined score by share of training data. Rows marked † and ‡ are cited from other papers, as the source caption states tab. 3. Full per-metric values are in tab. 3, and per-run values are in app. B tab. 10.

Dialogue state tracking, joint goal accuracy (%), full training data. Source: tab. 4. Rows marked † use a fixed ontology of all domain-slot pairs.

Low-resource DST on MultiWOZ 2.0, joint goal accuracy, mean ± standard deviation over five runs. Source: tab. 5.

Effect of the plug-and-play formulation against cascaded generation. The model is T5-small fine-tuned on MultiWOZ 2.0 without multi-task pre-training §5.1. Latency is for one example on one V100 GPU §5.1. Speedup is relative to SOLOIST's 208.69 ms tab. 6. Source: tab. 6.

Intent classification on Banking77, accuracy (%). Source: tab. 7. Rows marked † and ‡ are cited from other works, not re-run tab. 7.

Pre-training with different annotations (PPTOD-small, MultiWOZ 2.0 and Banking77). Source: tab. 8. The first row is T5-small with no multi-task pre-training; the last row is PPTOD-small with all four annotations.

Human evaluation of 50 MultiWOZ 2.0 dialogue sessions, PPTOD-base against SOLOIST, five graders, three-point scale (0, 1, 2) §5.3. Source: tab. 9.

Dataset annotations and sizes for pre-training. Source: tab. 1.

Paper vs code

Official repository: awslabs/pptod at commit 9bae5eafed830e9a354711b70dc90e0f495b471c, Apache-2.0. The paper's first-page footnote links it p. 1. The repository was found through that link.

Consistent with the paper:

• Batch size 128 in all three pre-training launch scripts [code: Pretraining/pretraining_pptod_base.sh#L4-L6] [code: Pretraining/pretraining_pptod_small.sh#L4-L6] [code: Pretraining/pretraining_pptod_large.sh#L4-L6], matching §3.4.
• Learning rate 5e-5 and 10 epochs pretrain.py pretrain.py, matching §3.4.
• Target tokens are wrapped with start and end markers such as <sos_r> and <eos_r> tokenize_all_datasets.py. The paper does not mention these markers.
• The stored E2E inference results match the MultiWOZ 2.0 rows of Table 2. The base run's file name reads Inform 89.2, Success 79.4, BLEU 18.62, Combined 102.92 . The small and large files match their rows too. The DST files in DST/inference_result/ match the Table 4 2.0 column (53.37, 51.5, 53.89).
• Plug-and-play inference: the dialogue action and the response are both generated from the database result, not from each other [code: E2E_TOD/inference_utlis.py#L137] inference_utlis.py.

Discrepancies:

• Loss scale. Eq. 2 sums over target tokens eq. 2. The model returns the Hugging Face T5 loss, which averages over target tokens with padding masked out T5Model.py dataclass.py.
• TaskMaster POL. Table 1 marks TaskMaster as annotated for POL tab. 1, but the code's mapping turns POL off for it [code: Pretraining/dataclass.py#L26].
• Maximum length. The paper says samples are truncated to 1024 §3.4. The code keeps only the last 900 context tokens [code: Pretraining/dataclass.py#L128] dataclass.py and caps targets at 128 tokens dataclass.py.
• Data mixing. Algorithm 1 shuffles data from different tasks together [alg. 1]. The default is session-level: sessions are shuffled, and all of a session's turns and tasks stay in sequence dataclass.py. A turn-level option exists dataclass.py.
• Optimizer and schedule. The paper says Adam with no schedule given §3.4. The code uses AdamW with zero weight decay and a linear decay schedule with no warmup pretrain.py [code: Pretraining/pretrain.py#L105].
• Fine-tuning optimizer. The paper does not state it. The released fine-tuning code uses Adafactor at lr 1e-3, not the pre-training Adam at 5e-5 learn.py learn.py learn.py.
• Dropout. Not stated in the paper. The code uses 0.2 for pre-training pretrain.py and 0.1 for E2E fine-tuning [code: E2E_TOD/learn.py#L35].
• Dev data. Pre-training evaluates on the test split files of each corpus dataclass.py. The intent-classification dev set is a random 10% of pooled CLINC, ATIS and SNIPS data, not an official split process_intent_classification_dataset.py.
• Low-resource runs. The paper reports five seeds per setting §4.1.4. The code sets no seed, and the released scripts cover only full training (*_full_training.sh). The --train_data_ratio option exists [code: E2E_TOD/learn.py#L27].
• Generation caps. Not stated in the paper. The code uses 120 tokens for belief state and response, and 80 for the dialogue action [code: E2E_TOD/inference_utlis.py#L79] inference_utlis.py inference_utlis.py.
• Dev score. The E2E checkpoint is chosen by 0.5 × (Success + Inform) + BLEU on the dev set learn.py.
• Latency. The paper's latency figures come from a V100 with batch size 1 §5.1. The repository lists no latency benchmark script. Not verified further.

Pitfalls

• Loss normalisation. Eq. 2 is a sum over target tokens eq. 2, but the code's loss is a mean over the non-padding target tokens of each batch [code: Pretraining/modelling/T5Model.py#L24-L25] dataclass.py. The mean therefore depends on how many real target tokens a batch holds. With Adam, a constant rescaling of the loss matters less than with SGD, but it still changes the effect of gradient clipping. Choose one convention and state it.
• Gradient accumulation. The loss is not divided by the accumulation count, and gradient clipping runs after every micro-batch, not only before each optimizer step pretrain.py. The fine-tuning scripts follow the same pattern learn.py. As written, the gradient seen by the clip is the sum over micro-batches. An implementer who copies this pattern gets a different effective step size from one that averages.
• Algorithm versus default shuffle. Algorithm 1 reads as mixing tasks [alg. 1]. The default keeps each session's turns and tasks adjacent [code: Pretraining/dataclass.py#L207-L214]. A 128-sample batch therefore contains mostly samples from a few sessions.
• Corpus count. The text says eleven corpora §3.1 but lists ten by name, leaving out CamRest676 §3.1. Appendix A and Table 1 include it app. A [tab. 1].
• Domain count. The domain column of Table 1 sums to 99. The text says 80 domains §3.1. The paper does not say whether domains are shared across corpora.
• TaskMaster POL. See the discrepancy above. A reimplementation that follows Table 1 and one that follows the code will train on different data.
• Fine-tuning hyperparameters. The paper gives none beyond "same objective" §3.3. The released Adafactor settings are the only values available.
• Low-resource seeds and splits. The five-run averages cannot be reproduced exactly from the code, because no seed is set and the low-resource launch scripts are absent. For the smallest ratios, the code skips evaluation for the first epochs (10, or 3 with pretrained initialisation), so checkpoints are chosen only after that point learn.py.
• Intent classification is not the best at 10 samples. USE+ConveRT reaches 85.19 against PPTOD-large's 84.12 tab. 7. The paper calls PPTOD comparable to existing methods §4.3, and at 10 samples it trails USE+ConveRT.
• Ontology-based DST rows. The best DST values in the 2.1 column come from ontology-based classification (FPDSC-dual, 59.07) tab. 4. PPTOD is generative and needs no fixed ontology §4.2.1.
• Re-run and cited numbers. Several table entries are cited from other papers (†/‡ in Tables 2, 3, 7). The UBAR scores come from author-released models tab. 2.
• Human evaluation scope. Only PPTOD-base versus SOLOIST, on 50 sessions, with five graders §5.3. Fleiss' κ for the truthfulness rating is 0.598 [tab. 9], so that rating is the least consistent across graders.
• Latency is from a single setup. Speedups in Table 6 are one V100 at batch size 1, with no variance reported §5.1.
• Variance is reported unevenly. Tables 5 and 7 give means and standard deviations over runs tab. 5 tab. 7. The paper does not say how many runs Tables 2 and 4 use §4.1.3 §4.2.1.

Lineage

• Builds on: T5 (Raffel et al., JMLR 2020; cited in the paper as DBLP:journals/jmlr/RaffelSRLNMZLL20) for the initial weights and text-to-text format §3.4. No arXiv id is given in the source, and no wiki card exists.
• Multi-task pre-training on dialogue corpora builds on earlier work the paper cites for unifying NLP tasks into one text format §3.2.
• Compared against: SOLOIST (Peng et al., 2021), TOP and TOP+NOD (Liu et al.), SimpleTOD, UBAR, MinTL, SOM-DST, DSTQA, Seq2seq-DU, TripPy and others §4.1.2 tab. 2 tab. 4. None of these has a card in this wiki yet.
• Superseded by: not known from wiki content.