jumafernandez/trace-deep-bidi

TRACE (deep, bidirectional)

TRACE is a contextual turn encoder for task-oriented dialogue. A frozen base encoder maps each turn to a static embedding $e_t$, and TRACE, a Transformer whose tokens are turns, maps the sequence $(e_1, \dots, e_T)$ to contextual representations $h_t$. In the bidirectional mode each turn attends to every turn of the window, before and after it, so $h_t$ describes the turn within the whole dialogue. Training is self-supervised and uses no dialogue act annotations.

The variant used by default is jumafernandez/trace-deep-ar. The checkpoints of every variant in the format of the training code are in jumafernandez/trace-checkpoints.

Model details

Base turn encoder (frozen)	sergioburdisso/dialog2flow-joint-bert-base
Architecture	BERT-style encoder over turn embeddings: 12 layers, 12 heads, hidden size 768
Parameters	85.1M
Attention	bidirectional
Inputs	turn embeddings (768), speaker ids and turn positions, up to 64 turns per window
Output	$h_t \in \mathbb{R}^{768}$ for every turn
Training data	task-oriented dialogue collection of Dialog2Flow, partitioned by near-duplicate clusters
Training	from scratch, self-supervised masked turn reconstruction, learning rate $10^{-4}$, checkpoint selected by validation loss
Seeds	42

Usage

The model ships its own code, so it loads with trust_remote_code=True. Pinning a revision is recommended.

from transformers import AutoModel

model = AutoModel.from_pretrained("jumafernandez/trace-deep-bidi", 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 = model.encode_dialogue(turns, speakers)   # (3, 768), one h_t per turn

encode_dialogue computes the base embeddings with the base encoder declared in the configuration, which requires sentence-transformers. Precomputed base embeddings can be passed with embeddings=. Dialogues longer than 64 turns are encoded in consecutive windows.

Since every turn attends to the turns after it, $h_t$ depends on the rest of the window. To obtain representations that use only the dialogue so far, each prefix is encoded and the representation of its last turn is kept:

e = model.base_embeddings(turns)
h_so_far = [model.encode_dialogue(turns[:t + 1], speakers[:t + 1], embeddings=e[:t + 1])[-1] for t in range(len(turns))]

This repository contains the checkpoint converted to the transformers format. The conversion was verified against the representations computed with the training code on the same base embeddings of 50 SpokenWOZ dialogues (883 turns), with a maximum absolute difference of $3.8 \times 10^{-6}$, and on every prefix of 20 of those dialogues (413 turns), with a maximum absolute difference of $4.5 \times 10^{-6}$.

Limitations

The model was trained on English task-oriented dialogue and expects embeddings of the base encoder it was trained on.

Citation

Fernández, J. M., Errecalde, M., & Burdisso, S. TRACE: Learning the Update for Contextual Turn Representations in Task-Oriented Dialogue. In preparation.