Exploring Embedding Priors in Prompt-Tuning for Improved Interpretability and Control

Prompt-Tuning is an efficient method for adapting pre-trained language models to new tasks with minimal computational overhead by modifying prompt embeddings. In this work, we investigate how crucial the phenomenon of embedding collapse, frequently observed in Prompt-Tuning, is for the final performance of the model. To address this question, we designed embedding priors and compared them with posteriors of the converged Soft and Deep Prompt-Tuning methods. Our findings suggest that priors strongly affect the position of the tuned embeddings, and models can effectively work with embeddings from different parts of activation spaces, including completely new regions. As the final Prompt-Tuning capabilities are limited, we hypothesize that controllable Prompt-Tuning posteriors may serve as a good starting point for tasks such as chain-of-thought (COT) distillation. Our experiments also show that generated trajectories are not localized in the activation space of the models. However, there are distinct clusters of activations for distant tasks (e.g., NLP and arithmetic), while activations between NLP tasks (e.g., Question-Answering and MLM) lie in the same cluster. These observations raise questions about the importance of a single activation cluster for the generalization abilities of large language models.

💡 Research Summary

Prompt‑tuning has emerged as a lightweight adaptation technique for large pre‑trained language models, requiring only a small set of trainable prompt embeddings while keeping the backbone frozen. A recurring observation in this setting is “embedding collapse”: the newly learned prompt vectors tend to converge toward existing token embeddings, reducing diversity and potentially limiting generalization. This paper asks how essential that collapse is for downstream performance and whether we can deliberately steer the embeddings away from the collapse using Bayesian priors.

The authors conduct a systematic study on LLaMA‑1B (16 layers) using two very different downstream tasks: the Stanford Question‑Answering Dataset (SQuAD) for natural‑language understanding and the DeepMind Math dataset for arithmetic reasoning. They evaluate two prompt‑tuning variants: (1) Soft Prompt, where 20 token‑level embeddings are prepended to the input, and (2) Deep Prompt, where an additional 20 trainable embeddings are inserted at the last three layers (plus the same 20 token embeddings). All model weights remain frozen.

Four families of priors are explored:

1. Isotropic Gaussian (N(0, σ²I)) – a simple, unstructured baseline.
2. Structured Gaussian (N(μ, Σ)) – fitted to the distribution of the original token embeddings, thus preserving inter‑dimensional correlations.
3. Gaussian Exclusion – a widened Gaussian (c·dim·Σ) from which samples are accepted only if their density under the original fitted Gaussian is low, explicitly pushing embeddings into low‑density regions of the original space.
4. Gaussian Interpolation – linear interpolation between samples drawn from the fitted prior (pre‑training domain) and a prior estimated on the target domain, intended to “bridge” the two activation clusters.

Additionally, a VAE‑based sampler is tested to generate embeddings from a smooth latent distribution that may better capture non‑linear structure across domains.

Visualization experiments (t‑SNE on token embeddings, PCA on high‑layer activations) reveal that sentence‑level trajectories are highly non‑local in both token and activation spaces; they do not form tight clusters even when the model processes the same input repeatedly. However, when comparing domains, the activations for SQuAD and Math form distinct clusters, while different NLP tasks (e.g., QA vs. masked‑language‑modeling) share a common cluster. This suggests that the model’s internal representation space is organized by task family rather than by individual examples.

The core of the analysis examines how the choice of prior and learning rate affect the final prompt embeddings. With an isotropic Gaussian initialization and a relatively high learning rate (5e‑3), prompt vectors tend to collapse toward the pre‑trained token cloud. Lowering the learning rate (5e‑4) or using a structured Gaussian initialization yields embeddings that diverge substantially from the original token distribution. Crucially, despite these large positional shifts, downstream performance (accuracy, precision, recall, F1) remains virtually unchanged across all prior configurations. Even the aggressive Gaussian‑exclusion prior, which forces embeddings into completely novel regions of the activation space, does not degrade task scores.

These findings lead to several key insights:

• Priors strongly dictate where prompt embeddings end up, but the location does not dictate performance for the tasks examined.
• Embedding collapse is not a strict bottleneck; the model can operate effectively with prompts situated in distant or even previously unseen activation zones.
• Domain‑specific activation clusters exist, implying that multi‑task or multi‑modality scenarios may benefit from priors that explicitly bridge clusters (e.g., Gaussian interpolation).
• Controllable posteriors could serve as useful initializations for more complex procedures such as chain‑of‑thought (CoT) distillation, where a well‑shaped prompt space might accelerate the learning of multi‑step reasoning.

The paper concludes that while embedding priors are a powerful tool for shaping the geometry of prompt embeddings, they do not, by themselves, improve raw task performance. Their true value lies in offering interpretability, enabling systematic exploration of the activation landscape, and providing a principled starting point for downstream methods that require richer, more controllable prompt representations. Future work may investigate how to exploit these priors for cross‑domain generalization, multi‑modal integration, and the systematic distillation of reasoning processes into compact prompt vectors.