Learning How to Ask: Querying LMs with Mixtures of Soft Prompts

In plain words

Masked language models such as BERT can recall facts when the question is written as a sentence with a blank, for example a sentence that asks for the date of a person's death. Until now, people wrote these sentences by hand or collected them from text, and the exact wording made a large difference.

This paper treats the sentence as something to learn. The words are replaced by continuous vectors that do not need to match any real word, and gradient descent adjusts them using known question and answer pairs. Several such learned sentences are combined, and the training also learns how much weight each one gets. On the benchmarks in the paper, this clearly beat the earlier prompt methods. Starting from random vectors worked almost as well as starting from a good hand-written sentence. The authors take this as a sign that the models hold more factual knowledge than the usual prompts reveal. The learned sentences are no longer readable as language, though.

TL;DR

Soft prompts are continuous token vectors plus a trainable additive perturbation at every layer, and a mixture of such prompts is fitted with EM for each relation §3.1–§3.5. The authors report P@1 gains on T-REx, Google-RE and ConceptNet over LAMA, LPAQA and AutoPrompt, using BERT and RoBERTa tab. 1 tab. 2 tab. 3. The main-text claim that the headline results tune only layer 0 conflicts with the numbers in Table 1, which match the all-layer variant §4.7 tab. 5. Before implementing, read Paper vs code and Pitfalls.

Problem

A fixed pretrained language model can be queried for a relation (for example date of death) with a cloze prompt that has two blanks: one for the subject x and one for the object y §3. Hand-written prompts are often brittle, and mined or paraphrased prompts (LPAQA) and gradient-searched discrete prompts (AutoPrompt) are limited to real words §2. The paper asks whether the prompt itself can be learned in continuous space so that it elicits the stored knowledge better §1.

Method

The task: given a fixed LM, a relation r and labelled pairs (x, y) for r, train a predictor of y from x, and test it on held-out pairs of the same relation §3. Only the prompt parameters are trained; the LM weights are frozen [code: soft_prompts/lm/language_model.py#L31-L33].

Soft prompts §3.1

1. Start from a hard prompt, a sentence with two blanks. Each of its tokens is replaced by a vector in R^d. The vectors are initialised to the embeddings of the hard prompt's tokens and are then tuned freely, so they need not match any word §3.1.
2. The number of tokens and their positions stay fixed §3.1.
3. For a six-token prompt the search space is 6d, the vector size times the token count §3.1.

Deeply perturbed prompts §3.2

4. Let v_i^(l) be the hidden vector of token i after layer l, with v_i^(0) the input embedding. Add a small trainable perturbation Δ_i^(l) to each token at each layer, before it feeds the next layer §3.2.
5. Initialise every Δ to zero §3.2.
6. Perturbing only layer 0 is the same as tuning the input vectors directly. Tuning all layers gives 6d(L+1) parameters for a six-token prompt, with L the layer count §3.2.
7. The authors expect small perturbations to keep activations close to what the LM saw in training. They keep them small with early stopping or another regulariser §3.2.

Mixture of prompts §3.3

8. For a relation r with prompt set P_r, the predicted distribution over y is a weighted sum over prompts of the LM's probability for y eq. 1 §3.3:

   p(y | x, r) = Σ_{P ∈ P_r} p(P | r) · p_LM(y | P, x) eq. 1

9. The weights p(P | r) form a distribution over the prompts. Product-of-experts ensembling is mentioned as an alternative §3.3.

Data-dependent weights, optional §3.4

10. The paper proposes making the weights depend on x through Bayes' rule, with a temperature T that is also learned. The authors report that this had little effect; training pushed T high enough to discard x §3.4 §4.7.

Training objective §3.5 §4.5

11. The objective is the negative log-likelihood of the correct y under the mixture, summed over training pairs eq. 2 §3.5:

    −Σ_{(x,y)} log Σ_{P ∈ P_r} p(y | P, x) eq. 2

    As printed, eq. 2 omits the mixture weights of eq. 1. The text says the weights are optimised jointly with the prompts §3.5. Use the weighted form.

12. The paper minimises this with EM. The E step computes, for each (x, y) pair, the posterior over prompts. The M step runs gradient descent on the prompt parameters, and the mixture weights are learned as a distribution §3.5. The released code trains the per-layer perturbations Δ and keeps the token vectors frozen relaxed_pattern.py.

13. Optimiser: Adam with its default configuration. Batch size 64 for gradient training. Early-stopping patience 4. Training runs up to 16 epochs, and the checkpoint with the best dev score is kept §4.5.

14. Predictions are the most probable y under the mixture §4.6. The answer y is always a single token §4.2.

# Build one soft prompt from a hard prompt  [§3.1] [§3.2]
# Input:  hard prompt tokens; LM depth L; vector size d
# Output: token vectors and per-layer perturbations

def make_soft_prompt(hard_tokens, num_layers):
    # Initialise each token from its embedding  [§3.1]
    vectors = embed(hard_tokens)                 # (T, d)
    # Zero perturbation for layers 0..L  [§3.2]
    # [code: soft_prompts/components/pattern.py#L303]
    bias = zeros(num_layers + 1, len(hard_tokens), d)  # (L+1, T, d)
    return vectors, bias

# Mixture prediction for one candidate answer  [§3.3] [eq. 1]
# Input:  prompts with weights; input x; candidate y

def mixture_probability(prompts, weights, x, y):
    # Weighted sum of the LM probability of y under each prompt  [eq. 1]
    total = 0.0
    for prompt, weight in zip(prompts, weights):
        total += weight * lm_probability(prompt, x, y)
    return total

# EM training for one relation  [§3.5] [§4.5]
# Input:  train pairs; initial prompts; dev pairs
# Output: best prompt perturbations and weights, by dev P@1

def train_mixture(train_pairs, prompts, dev_pairs, max_epochs=16, patience=4):
    # Start with uniform mixture weights  [§4.7]
    # [code: soft_prompts/model/relaxed_pattern.py#L43]
    weights = uniform(len(prompts))
    best_score = 0.0
    best_state = None
    for epoch in range(max_epochs):
        # E step: posterior over prompts for every pair  [§3.5]
        # [code: soft_prompts/model/relaxed_pattern.py#L150-L168]
        resp = responsibilities(train_pairs, prompts, weights)  # (N, K)
        # M step, weights: normalised responsibility per prompt
        # [code: soft_prompts/model/relaxed_pattern.py#L204-L209]
        weights = resp.sum(axis=0) / resp.sum()
        # M step, perturbations: Adam over shuffled (pair, prompt) batches
        # [§4.5] [code: soft_prompts/trainer/relax_pattern_trainer.py#L127-L138]
        for batch in shuffled_pairs(train_pairs, prompts):
            # Responsibility-weighted cross-entropy, summed over the batch
            # [code: soft_prompts/lm/language_model.py#L121-L127]
            loss = weighted_cross_entropy(batch, resp)
            adam_step(loss)
        # Dev check and early stopping
        # [code: soft_prompts/trainer/relax_pattern_trainer.py#L261-L270]
        score = dev_precision_at_1(dev_pairs, prompts, weights)
        if score > best_score:
            best_score = score
            best_state = snapshot(prompts, weights)
            patience = patience_limit
        elif patience == 0:
            break
        else:
            patience -= 1
    return best_state

Hyperparameters

Results

Author-reported. Setup: P@1 unless stated. Baselines are from prior papers or the authors' reimplementations §4.6. Significance is marked in the paper with a sign test (p < 0.02) on T-REx tab. 1 and a permutation test in the appendix tab. 5.

T-REx, the main comparison tab. 1:

The AutoPrompt row with 43.3 / 45.6 sits in the BERT-base block without a model label, and the text does not say which LM those numbers use §4.6 tab. 1.

Google-RE with BERT-large-cased tab. 2, P@1 / P@10 / MRR:

ConceptNet tab. 3: the paper reports random initialisation as the winner. LAMA with BERT-large: P@1 0.1, P@10 5.0, MRR 1.9. Soft (ran.) with BERT-large: P@1 14.5, P@10 38.6, MRR 22.1. Soft (min.) with BERT-large: P@1 12.8, P@10 37.0, MRR 20.9 tab. 3.

Ablation on T-REx original with BERT-large tab. 4, P@1 / P@10 / MRR:

Relation split with distinct y values tab. 9, P@1 / P@10 / MRR: LPAQA with BERT-base 18.9 / 40.4 / 26.6 becomes soft 23.0 / 45.2 / 30.5. With BERT-large 23.8 / 47.7 / 32.2 becomes soft 27.0 / 51.7 / 35.4 tab. 9.

Per-layer detail for BERT-large on T-REx original, soft mixture (min.) tab. 5: initial 39.4, after tuning layer 0 only 51.0, after tuning all layers 51.6 for P@1 tab. 5.

Summary of claims: soft prompts beat earlier prompt methods even when the words of the informed prompt are replaced by random vectors §4.7. Random initialisation is close to informed initialisation: BERT-large on T-REx original gives Soft (ran.) 51.9 against Soft (min.) 51.6 tab. 1. The authors conclude that LMs know more than earlier prompt-based probes showed §1 §5. Tuned prompts are not interpretable. Nearest-word projections of the tuned vectors are mostly on one topic but do not form phrases app. C fig. 1.

Paper vs code

Official code: the paper's conclusion links github.com/hiaoxui/soft-prompts, at commit eb9047e134de25dd9442dc2a809652814a005b4d, license Apache-2.0. The README and config.yaml replicate T-REx extended with BERT-large-cased [code: README.md] config.yaml.

• Batch size. Paper: 64 for gradient training §4.5. Released config: batch_size_grad: 32 config.yaml. The code's batch counts (relation instance, prompt) pairs, shuffled over the product of training instances and prompts relaxed_pattern.py.
• Layers tuned. Paper's text: the basic method tunes only layer 0 and the all-layer variant is an add-on §4.7. Table 1 numbers match the all-layer column of Table 5, not the layer-0 column tab. 1 tab. 5. The released config tunes all layers (max_layer: 24) and has no switch for layer 0 only config.yaml. Setting max_layer to 0 would leave only the layer-0 perturbation, but this is inferred from reading the code, not run pattern.py modeling_bert.py.
• Mixture weights. Paper: weights are a learned distribution, optimised by gradient through a softmax parameterisation §3.5. Code: weights are a tensor without gradient, set in closed form at each M step as the normalised sum of responsibilities per prompt relaxed_pattern.py relaxed_pattern.py.
• Trainable parameters. Paper: the prompt token vectors and the perturbations are both tuned §3.1 §3.2. Code: only the perturbations (bias) are passed to the optimiser; the token vectors have gradients disabled relaxed_pattern.py pattern.py.
• Data-dependent weights. Paper: p(P | r, x) with temperature T, learned §3.4. Code: the flag conditional_prompt is off config.yaml. When on, the weights are multiplied by one vector, computed once per evaluation set by summing log-probabilities over every relation instance in that set. It is not per example and has no temperature relaxed_pattern.py relaxed_pattern.py.
• Loss. Paper: negative log-likelihood of the mixture eq. 2. Code: responsibility-weighted cross-entropy summed over the batch, not averaged language_model.py. This is the EM surrogate, not eq. 2 directly.
• Early stopping. Paper: patience 4 §4.5. Code stops at the fifth consecutive epoch without a dev gain relax_pattern_trainer.py.
• BART. The paper reports BART results tab. 5. The released code raises NotImplementedError for bart construct_lm.py. BART results cannot be reproduced from this code.
• Random init. Paper: Gaussian fitted to all LM embeddings §4.4. Code: each prompt vector is drawn per dimension from the embedding mean and standard deviation, so covariance is diagonal relaxed_pattern.py. The released config sets randomize_prompt: false config.yaml.
• Data split. Paper: T-REx split inherited from AutoPrompt, other datasets split 80/10/10 at random §4.3. Code: T-REx ships pre-split files db_loader.py. The random split exists in the trainer [code: relax_pattern_trainer.py#L46-L47], but no Google-RE or ConceptNet config is released.
• Undocumented options. The code has a frequency prior (frequent, off) and a generate training target (commented out in config). Neither is described in the paper config.yaml relax_pattern_trainer.py.
• Entity truncation. Code truncates entity tokens to 16 [code: soft_prompts/components/pattern.py#L355] pattern.py. The paper does not mention this.

Pitfalls

• Layer-0 versus all-layer ambiguity. The main text and Table 1 disagree §4.7 tab. 1 tab. 5. Implement all layers, which is what the release does, and report which variant you run. Layer-0-only gives roughly 0.6 P@1 points less in the BERT-large example tab. 5.
• Eq. 2 omits weights. Implement the weighted mixture of eq. 1 eq. 1 eq. 2.
• Learning rate and loss scale not stated. The paper gives no learning rate. The code passes none and uses the framework default. The loss is summed over the batch, so its scale grows with batch size. The paper does not say how this interacts with the optimiser language_model.py.
• Batch size conflict. 64 in the paper, 32 in the release §4.5 config.yaml.
• Early-stopping off-by-one. Five stagnant epochs stop training, not four relax_pattern_trainer.py.
• Model selection. The reported test scores come from the checkpoint with the best dev P@1 relax_pattern_trainer.py. Keep dev and test separate when you tune anything else.
• Data-dependent weights are off in the release and not what Section 3.4 describes relaxed_pattern.py.
• Perturbation size is not regularised by default (penalty zero) config.yaml. The paper's claim that perturbations stay small depends on early stopping §3.2.
• Entity length cap of 16 tokens pattern.py.
• Environment. The release pins Python 3.7, torch 1.7.0 and transformers 4.0.0 requirements.txt README.md. It ships its own modified BERT, so newer library versions may not run it unchanged modeling_bert.py.
• Cost. A single prompt trains in minutes even for BERT-large, and mixture training scales roughly linearly in the number of prompts §4.5. Each epoch runs an E step over the whole training set before and after the update relax_pattern_trainer.py relax_pattern_trainer.py.
• Results are author-reported. The baselines mix paper numbers and reimplementations §4.6 tab. 1.
• Learned prompts are not readable. Use them as they are; they are not interpretable §5 app. C.

Lineage

• Builds on: LAMA (Petroni et al., 2019) for cloze-style knowledge probing §2; LPAQA (Jiang et al., 2020) for mined and paraphrased prompts §2; AutoPrompt (Shin et al., 2020) for gradient-guided discrete prompts §2.
• Compared against: LAMA, LPAQA, AutoPrompt on T-REx, Google-RE and ConceptNet tab. 1 tab. 2 tab. 3.
• Related, unpublished at submission: prefix-style soft tokens for generation and gradient-tuned continuous prompts for GPT-2 and BERT §2.