jumafernandez/trace-dialog-deep-tod

TRACE-Dialog (deep, task-oriented, supervised)

TRACE-Dialog produces a single embedding of a whole dialogue. It extends the TRACE contextual turn encoder with a learned [CLS] vector, placed before the turns and followed by a learned [SEP] vector, and the representation of the dialogue is the output $h_{[\mathrm{CLS}]}$ at that position. The slot is first trained without an objective of its own, following RoBERTa, on dialogues packed without cross-contamination, and then with a supervised similarity stage in which dialogues that share the same set of services are positives. The representation is sensitive to the trajectory of the dialogue, that is, to the order of its turns.

Model details

Base turn encoder (frozen)	sergioburdisso/dialog2flow-joint-bert-base
Architecture	BERT-style encoder over turn embeddings with [CLS] and [SEP] vectors: 12 layers, 12 heads, hidden size 768
Parameters	85.1M
Attention	bidirectional
Inputs	turn embeddings (768) and speaker ids, up to 62 turns
Output	$h_{[\mathrm{CLS}]} \in \mathbb{R}^{768}$ for the dialogue
Training data	task-oriented dialogue collection of Dialog2Flow; similarity stage supervised with the service sets of SGD
Training (similarity stage)	5 epochs, learning rate $2 \times 10^{-5}$, temperature 0.05, seed 42

Usage

The model ships its own code, so it loads with trust_remote_code=True. Pinning a revision is recommended. In both cases the turns are encoded with the base encoder declared in the configuration, and only the first 62 turns of a dialogue are used.

With sentence-transformers

Each input is a whole dialogue, given as a list of chat messages (role and content) or of turns (speaker and text). The roles user and assistant (or system) are the user and the system of the dialogue.

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("jumafernandez/trace-dialog-deep-tod", trust_remote_code=True)

dialogue = [{"role": "user", "content": "hi, I need a cheap restaurant in the centre"},
            {"role": "assistant", "content": "sure, what kind of food would you like?"},
            {"role": "user", "content": "italian"}]
embeddings = model.encode([dialogue])   # (1, 768), one h_CLS per dialogue

The embeddings can be compared with the similarity functions of sentence-transformers, for example to retrieve the dialogues of a collection that are closest to an ongoing one.

With transformers

from transformers import AutoModel

model = AutoModel.from_pretrained("jumafernandez/trace-dialog-deep-tod", trust_remote_code=True)

turns = ["hi, I need a cheap restaurant in the centre", "sure, what kind of food would you like?", "italian"]
speakers = ["user", "system", "user"]
h_cls = model.encode_dialogue(turns, speakers)   # (768,)

encode_dialogue also accepts precomputed base embeddings with embeddings=.

Evaluation

In the retrieval of conversational situations on SpokenWOZ, where the next dialogue act is predicted by the vote of the continuations of the $k$ nearest turns (1,019 single-domain training dialogues, cross-validated), the dialogue representation alone summarizes the services of the conversation rather than its last turn, and it is most useful combined with the embedding of the current turn:

Representation	Accuracy	Macro-F1	Coverage@5
$e_t$ (Dialog2Flow)	0.454	0.135	0.646
$h_{[\mathrm{CLS}]}$ of the prefix (this model)	0.301	0.084	0.548
$[e_t ; \lambda h_{[\mathrm{CLS}]}]$	0.482	0.173	0.685

This combination is the representation of the conversational memory of the dialog2agent agent.

This repository contains the checkpoint converted to the transformers format. The conversion was verified against the representations computed with the training code on 1,882 dialogue prefixes of SpokenWOZ, with a maximum absolute difference of $7.4 \times 10^{-6}$. The sentence-transformers interface was verified from the text of 500 dialogue prefixes of SpokenWOZ, from 1 to 90 turns, with a maximum absolute difference of $1.4 \times 10^{-5}$ with sentence-transformers 5.6 and 6.2.

Limitations

The model was trained on English task-oriented dialogue and expects embeddings of the base encoder it was trained on. The similarity stage relies on the services of SGD, so the notion of similar dialogues follows the task they serve.

Citation

Fernández, J. M., Errecalde, M., & Burdisso, S. TRACE-Dialog: Trajectory-Aware Dialogue Embeddings from Contextual Turn Representations. In preparation.