Prompt Tuning › Prompt Length Analysis

prompt_lengths = [1, 5, 10, 20, 50, 100]
final_losses = []
convergence_speeds = []  # Epochs to reach loss < 0.3

for p_len in prompt_lengths:
    # Reset before construction so every prompt length starts from a
    # reproducible vocabulary sample and the same pretrained checkpoint.
    torch.manual_seed(42)
    base_model_temp = AutoModel.from_pretrained(model_name)
    model_temp = PromptTuningClassifier(
        pretrained_model=base_model_temp,
        tokenizer=tokenizer,
        num_prompt_tokens=p_len,
        num_classes=2,
        init_from_vocab=True,
    )

    # Reset the classifier separately so its initial weights do not depend on
    # how many random vocabulary rows the prompt initializer sampled.
    torch.manual_seed(43)
    model_temp.classifier.reset_parameters()
    optimizer = torch.optim.AdamW(
        model_temp.parameters(), lr=3e-4, weight_decay=0.01
    )

    model_temp.train()
    convergence_epoch = None
    epoch_loss = None

    for epoch in range(50):
        batch_losses = []
        for i in range(0, len(texts), batch_size):
            optimizer.zero_grad()
            logits = model_temp(
                input_ids[i : i + batch_size],
                attention_mask[i : i + batch_size],
            )
            loss = F.cross_entropy(logits, labels_tensor[i : i + batch_size])
            loss.backward()
            optimizer.step()
            batch_losses.append(loss.item())

        # Compare prompt lengths using the same full-dataset epoch loss as the
        # chapter's main training loop, rather than one arbitrary first batch.
        epoch_loss = float(np.mean(batch_losses))
        if convergence_epoch is None and epoch_loss < 0.3:
            convergence_epoch = epoch + 1

    final_losses.append(epoch_loss)
    convergence_speeds.append(
        convergence_epoch if convergence_epoch is not None else 50
    )

embedding_dim = 768
classifier_params = 768 * 2 + 2
# Use the base model loaded earlier to get exact count
base_model_params = sum(p.numel() for p in base_model.parameters())

efficiency_results = []
for p_len in prompt_lengths:
    prompt_params = p_len * embedding_dim
    total_trainable = prompt_params + classifier_params
    efficiency = (total_trainable / base_model_params) * 100
    efficiency_results.append((p_len, total_trainable, efficiency))