TSDAE: Using Transformer-based Sequential Denoising Auto-Encoder for Unsupervised Sentence Embedding Learning

In plain words

Good sentence embeddings usually need many labelled pairs of similar sentences, and labelling is expensive. This paper asks whether a sentence encoder can learn from plain, unlabelled sentences alone. The method is a game of reconstruction: the input is a sentence with many of its words removed, and the model has to rebuild the full sentence. The only route for that information is a single vector produced by the encoder, so that vector has to hold the meaning of the sentence.

After training, the decoder is thrown away and the encoder alone produces the embeddings. The authors also argue that the usual test, similarity on general news-style sentence pairs, says little about how a model will do on specialised text. They therefore test on search, re-ranking and paraphrase tasks drawn from different domains. On those tasks the method beats other unsupervised approaches, and it is also a good starting point for further training on labelled data.

TL;DR

• The encoder reads a sentence with 60% of its words deleted and produces one vector. A decoder rebuilds the original sentence from that vector only §3.1 §3.2 app. A.
• The decoder's cross-attention sees only the sentence vector, which forces the vector to carry the meaning §3.2.
• Author-reported: TSDAE has the best unsupervised average over four domain tasks, 55.2 against 52.4 for CT tab. 2. It is weak on general STS: 66.0 Spearman on STSb against 73.9 for CT tab. 3.
• Author-reported: as a pre-training step before labelled training, TSDAE beats masked language modelling on most datasets §6 fig. 2.
• Training runs a second Transformer, a decoder with cross-attention, next to the encoder denoising_auto_encoder.py. Only the encoder is used at inference §3.2.

Problem

Sentence embeddings for search, re-ranking and paraphrase detection are usually trained with labelled sentence pairs. For most domains those labels are not available, and making them is costly §1. Earlier unsupervised methods were mostly judged on STS, but STS sentences are general-domain, have an artificial score distribution, and do not reflect how the embeddings do on real tasks §4.

The paper asks how unsupervised methods perform on specialised domains, whether labelled data can be added later, and whether the method helps as pre-training §1 §4.

Method

Training objective. The encoder receives a damaged sentence x̃ and produces a fixed-size sentence vector. The decoder must predict every token of the original sentence x from that vector. The objective is the expected log-likelihood of x given x̃, summed over the tokens of x §3.1. The output probability of each token is a softmax over the vocabulary embeddings, scored against the decoder hidden state §3.1. The paper's displayed objective is unnumbered; the only numbered equations in the paper are the two BERT-flow equations eq. 1 eq. 2.

Damage. The noise is token deletion. The paper's text names deletion and swapping §3.1; its appendix compares five noise types and picks deletion app. A tab. 5. The deletion ratio is the fraction of tokens affected, and the search selected 0.6 app. A tab. 6. The paper's stated example is a ratio of 0.3 meaning 30% of tokens are affected app. A.

Decoder bottleneck. The paper's modified cross-attention takes its keys and values only from the sentence vector, as a one-row matrix of size d, so the decoder never sees the encoder's per-token states §3.2. The decoder hidden state at each decoding step attends to that single row §3.2.

Sentence vector. The sentence vector is the [CLS] token output, selected in the appendix search app. A tab. 7.

Weight tying. Encoder and decoder share parameters during training. Cross-attention weights have no encoder counterpart and are not shared app. A. The paper does not say how the decoder is initialised; the code starts it from the same checkpoint as the encoder when weights are tied denoising_auto_encoder.py.

Training. For TSDAE the paper uses 100K steps, batch size 8, AdamW, learning rate 3e-5 app. B. It does not use early stopping on labelled data, because labelled data is assumed to be unavailable app. B.

# One TSDAE training step  [§3.1] [§3.2]
# Code citations are relative to sentence_transformers/sentence_transformer/
# Input:  original sentences, token ids (B, L)
# Output: scalar loss

def tsdae_loss(encoder, decoder, original_ids, del_ratio=0.6):
    # Damage each sentence by deleting tokens  [app. A] [tab. 6]
    damaged_ids = delete_tokens(original_ids, del_ratio)
    # One sentence vector per input, from the [CLS] position  [§3.2] [app. A]
    sentence_vecs = encoder(damaged_ids)             # (B, d)
    # The decoder sees only this vector, as a memory of length 1  [§3.2]
    memory = sentence_vecs[:, None, :]               # (B, 1, d)
    # Teacher forcing: predict each next token of the original  [§3.1]
    decoder_input = original_ids[:, :-1]             # (B, L-1)
    labels = original_ids[:, 1:]                     # (B, L-1)
    # Token logits from the decoder, attending to the memory  [§3.2]
    logits = decoder(decoder_input, memory)          # (B, L-1, V)
    # Average cross-entropy over non-padding tokens
    # [code: losses/denoising_auto_encoder.py#L213-L214]
    loss = cross_entropy(logits, labels, ignore=PAD_ID)
    return loss


# Random deletion of words  [app. A]
# [code: datasets/denoising_auto_encoder.py#L49-L61]
def delete_tokens(sentence_ids, ratio):
    # Split the sentence into words  [code: datasets/denoising_auto_encoder.py#L53]
    words = to_words(sentence_ids)
    # Drop each word with probability ratio  [code: datasets/denoising_auto_encoder.py#L58]
    keep = [random() >= ratio for word in words]
    # If every word was dropped, keep one random word
    # [code: datasets/denoising_auto_encoder.py#L59-L60]
    if not any(keep):
        keep[random_index(len(words))] = True
    return to_ids([w for w, k in zip(words, keep) if k])

Steps:

1. Sample a batch of sentences and damage each one by deleting tokens at ratio 0.6 app. A tab. 6.
2. Encode the damaged sentence and take the CLS vector as the sentence embedding §3.2 app. A.
3. Feed the vector to the decoder as the only encoder-side input and train the decoder to output the original tokens, shifted by one position §3.1 §3.2.
4. Update encoder and decoder together, since their weights are shared app. A app. B.
5. After training, use the encoder alone with the pooling in step 2 to embed sentences §3.2.

Hyperparameters

Results

Author-reported. Setup: BERT-base-uncased with mean pooling unless noted; four domain datasets (AskUbuntu, CQADupStack, TwitterPara, SciDocs), each scored by MAP or AP and averaged across sub-tasks; 5 seeds §4.1 §5 tab. 2.

Main table tab. 2. The paper's † marks cases where TSDAE beats both CT and SimCSE in all five runs. The Avg. column is the mean of the four dataset scores.

USE-large was trained on in-domain data for AskUbuntu and CQADupStack; those values are shown in parentheses tab. 2. The per-dataset MAP or AP scores are as the paper reports them; the paper gives no per-task variance for this table.

Checkpoint study tab. 9, mean ± standard deviation over 5 seeds. Each checkpoint gets the four unsupervised methods that tab. 9 reports; there are no MLM rows.

The paper marks best values within each group by underline and the overall best by bold tab. 9; those marks are not reproduced here. Each value is shown as the paper gives it.

Other STS and pre-training numbers:

Table 3 gives the averages; the per-dataset STS table is tab. 11. The label-efficiency numbers in tab. 10 are estimates from a binary search over labelled-pair counts with precision set to one standard deviation of the score app. D. The annotation-time estimate, 20 minutes per 100 Twitter pairs, is from one annotator in one experiment app. D.

Ablations app. A tab. 5 tab. 6 tab. 7. Noise type, ratio and pooling, scored by STS-B Spearman on the tuning run:

Whether these ablation scores are on the validation or test split is not stated; the text calls the tied-weight result a validation score app. A.

Paper vs code

The paper's footnote links UKPLab/sentence-transformers. The code cited here is the repository at commit 1fa45b4 (version 6.2.0.dev0 in pyproject.toml), which is much newer than the v0.3.8 the paper says it used app. B pyproject.toml. The clone is shallow and contains no v0.3.8 code, so the code below is the current implementation. Whether it reproduces the paper's tables was not checked.

• Deletion probability, example scripts. The paper's ratio 0.6 means 60% of tokens are affected app. A. The legacy dataset deletes each word with probability 0.6 denoising_auto_encoder.py and the README keeps each word with probability 1 − 0.6 README.md. The STS-B example script keeps each word with probability 0.6, so it deletes about 40% train_stsb_tsdae.py. The file-based example uses the same rule train_tsdae_from_file.py. Neither matches the paper's 60% deletion.
• Loss normalisation. The paper's objective sums log-probabilities over tokens for each sentence §3.1. The code averages cross-entropy over non-padding tokens in the batch denoising_auto_encoder.py. The learning rate tuned under one normalisation may not transfer to the other.
• Decoder causality. The decoder is loaded as a causal language model with cross-attention enabled denoising_auto_encoder.py. The paper does not state whether the decoder self-attention is causal.
• Decoder input and target. Teacher forcing feeds the original sentence without its last token and predicts it shifted by one denoising_auto_encoder.py. The paper gives no shift rule.
• Weight tying direction. The tying replaces encoder parameters with the decoder's objects, and the cross-attention weights are skipped denoising_auto_encoder.py denoising_auto_encoder.py. The paper does not describe direction or the cross-attention exception.
• Sentence vector. The example scripts use CLS pooling train_stsb_tsdae.py, which matches app. A. The README snippet does not set pooling, so it uses the library default.
• Training length. The paper trains 100K steps at batch 8 app. B. The STS-B example trains for one epoch over the wiki1m dataset at batch 8, which gives about 124K steps for the 990K training rows shown in its comment train_stsb_tsdae.py train_stsb_tsdae.py train_stsb_tsdae.py.
• Mixed precision and sequence length. The example uses fp16 and a maximum sequence length of 75 tokens train_stsb_tsdae.py train_stsb_tsdae.py. The paper states neither.
• Static embeddings are refused. The loss raises an error for models whose first module is a static embedding denoising_auto_encoder.py.

No discrepancies were checked for the baselines (CT, SimCSE, BERT-flow, MLM), because their code is not in this repository. The paper says it re-implemented BERT-flow and used the official SimCSE code §5.5 app. B.

Pitfalls

• Score-table mismatch. The Table 2 caption says "average precision", but Section 4 says MAP for AskUbuntu, CQADupStack and SciDocs and AP for TwitterPara §4.1 tab. 2.
• Gain over CT. The text says 2.6 points on average over all tasks §6. Table 2 gives 55.2 against 52.4, a 2.8-point gap tab. 2. The "up to 6.4 points" claim is the SciDocs average tab. 2.
• "93.1% of supervised". The abstract's 93.1% is the AskUbuntu ratio, 59.4 over 63.8 for in-domain SBERT tab. 2. The other datasets give lower ratios, for example TwitterPara 73.0 over 78.7.
• Pooling choice. Mean pooling scored slightly higher than CLS (78.84 against 78.77), but the paper chose CLS because mean pooling loses position information app. A tab. 7. An implementer who wants the best score on the ablation can take mean pooling.
• Decoder initialisation. When tying is on, the decoder is initialised from the encoder checkpoint, and its cross-attention weights are new and random. The paper does not say how cross-attention is initialised denoising_auto_encoder.py train_stsb_tsdae.py.
• Deletion versus swap. The text names deletion and swapping as noise §3.1, but the code only implements deletion denoising_auto_encoder.py. The appendix comparison of five noise types is the only evidence for other types app. A tab. 5.
• Ablation split. The tuning run trains on SNLI+MultiNLI and reports STS-B Spearman app. A. Whether those scores are validation or test is not stated. The noise type, ratio and pooling were chosen on these STS-B scores, so the STS-B numbers in the main tables are not a clean held-out result for that configuration.
• Dev-set selection. The main experiments do not early-stop on labelled data app. B. The corpus-size study reports the best development score across epochs §7.1. That is a different protocol from the main table.
• STS does not predict downstream. TSDAE trails CT by 7.9 points on STSb tab. 3. The paper argues that STS is a weak selection criterion for the downstream tasks it reports §4 §8 §9; the implementer should not expect STS gains from TSDAE.
• Cost. Each training step runs the encoder and the decoder, and the code's docstring says tying the weights limits memory use denoising_auto_encoder.py. The paper does not report training time or memory.
• Decoder checkpoint. Using BART or T5 as the base made the decoder overfit: low training loss but worse test scores than training from scratch §8 tab. 4. The paper recommends single-encoder checkpoints such as BERT.
• Evaluation toolkit. The paper's evaluation splits and toolkit are at UKPLab/useb §4.1. That repository was not fetched or checked.
• Figures and equations. The paper's figures are described in words only here; they are not reproduced. Only two equations are numbered in the paper eq. 1 eq. 2; the TSDAE objective is cited by section §3.1.
• Commented-out content. The TeX source includes commented-out sections (dataset statistics, training-dynamics figures, an MNRL section). They are not part of the compiled paper and are not used here.
• Sentence sampling for STS. The STS-analysis training corpus is "sampled from Wikipedia" §6.1. Its size and sampling procedure are not stated.

Lineage

• Builds on: Sequential denoising autoencoders, cited as Vincent et al. 2010 and Goodfellow et al. 2016 §3; the Transformer encoder-decoder architecture §3.1, changed so the decoder reads only the sentence vector §3.2.
• Compared against: Contrastive Tension, cited as Carlsson et al. 2021 §5.3; SimCSE, cited as Gao et al. 2021 §5.4; BERT-flow, cited as Li et al. 2020 §5.5; masked language modelling as in BERT §5.2; Sentence-BERT in its supervised and out-of-the-box forms §5.1; USE-large §5.1.
• Superseded by: not stated.