The Geometry of Prompting: Unveiling Distinct Mechanisms of Task Adaptation in Language Models (2502.08009v1)

Abstract: Decoder-only LLMs have the ability to dynamically switch between various computational tasks based on input prompts. Despite many successful applications of prompting, there is very limited understanding of the internal mechanism behind such flexibility. In this work, we investigate how different prompting methods affect the geometry of representations in these models. Employing a framework grounded in statistical physics, we reveal that various prompting techniques, while achieving similar performance, operate through distinct representational mechanisms for task adaptation. Our analysis highlights the critical role of input distribution samples and label semantics in few-shot in-context learning. We also demonstrate evidence of synergistic and interfering interactions between different tasks on the representational level. Our work contributes to the theoretical understanding of LLMs and lays the groundwork for developing more effective, representation-aware prompting strategies.

• The paper analyzes how instruction, demonstration, and soft prompts affect the geometry of language model representations, revealing distinct underlying mechanisms despite similar task performance.
• Using manifold capacity and geometric analysis, the study highlights the critical roles of input examples and label semantics in few-shot in-context learning, showing how demonstrations reorganize intermediate layers.
• Findings suggest performance variations are often due to readout misalignment rather than representational quality, indicating potential for representation-aware strategies like prompt-tuning to improve alignment.

The paper "The Geometry of Prompting: Unveiling Distinct Mechanisms of Task Adaptation in LLMs" investigates how different prompting methods affect the geometry of representations in decoder-only LLMs. It employs a framework grounded in statistical physics and manifold capacity to reveal that various prompting techniques operate through distinct representational mechanisms for task adaptation, despite achieving similar performance. The paper highlights the critical role of input distribution samples and label semantics in few-shot ICL (In-Context Learning). The authors demonstrate evidence of synergistic and interfering interactions between different tasks on the representational level, contributing to the theoretical understanding of LLMs and laying the groundwork for more effective, representation-aware prompting strategies.

The paper explores the underlying mechanisms of ICL (In-Context Learning) by analyzing how different prompting methods modify internal representations in pre-trained LLMs. The authors analyze the separability and geometric properties of category manifolds, which are point clouds in the model's embedding space that correspond to examples sharing a category label. They use the framework of manifold capacity, which analytically connects task performance to the geometric properties of these representations, to illuminate ICL.

The core contributions include:

• A comprehensive analysis of how various prompting methods affect internal representations in LLMs, revealing distinct computational mechanisms despite similar performance outcomes.
• Novel insights into ICL (In-Context Learning) dynamics, including the role of label semantics, synergistic effects of demonstrations on unrelated tasks, and representational trade-offs during task adaptation.

The paper distinguishes between conventional few-shot ICL (In-Context Learning) and other input-based task-adaptation methods, referring to it as providing demonstrations to emphasize the role of task examples. ICL (In-Context Learning) performance depends heavily on the exact choice of examples, their ordering, and formatting. The actual input-output mapping matters less than expected, suggesting that few-shot ICL (In-Context Learning) involves an interplay of true task learning from examples and task recognition from the pre-training corpus. The authors distinguish between demonstration prompts and instruction prompts, noting that these two prompt types affect internal representations differently, despite comparable performance. Prompt-tuning is presented as an alternative approach to task adaptation, where a small set of continuous vectors (soft prompts) are learned and concatenated to the input embeddings while the model parameters remain frozen.

The paper leverages the emerging linear representation hypothesis, which posits that the embedding space contains "feature directions" encoding human-interpretable concepts, enabling the model to perform vector operations with meaningful semantics. The concept of feature superposition is discussed as a means by which a model can operate on more features than it has orthogonal directions in the embedding space. The authors note a lack of understanding of how the context preceding a given input (particularly, task adaptation) affects feature representation, given that probing methods are usually performed on a very diverse input corpus of text and averaged.

The paper discusses the concept of manifold untangling, where collective representations of inputs sharing a target category (a category manifold) must be well-separated from other categories. The framework of manifold capacity is presented as a formal link between representational geometry and separability. Manifold capacity quantifies how efficiently task-relevant features are encoded from the perspective of a linear downstream decoder, measuring the separability of target classes in the embedding space.

To investigate the effects of various prompting methods on representations in different task-specific contexts, the authors generated a synthetic dataset tailored to their research requirements using a LLM (Claude 3.5 Sonnet). The synthetic dataset consists of diverse sentences, each simultaneously labeled with three types of categories: Sentiment, Topic, and Intent, with five categories for each type. All experiments, including those focused on single-task performance, utilized this dataset, with the sentiment classification task serving as the primary focus. Key single-task experiments were replicated using established open datasets as a control.

The work focuses on text classification tasks with a fixed set of categories, which provides an analytically-grounded link between the geometry of the underlying representation and the separability of categories in the embedding space. In decoder-only LLMs, performance is affected by the representation quality and the readout alignment (the alignment between the model's unembed layer and the ideal decoder directions). Manifold capacity theory allows the disentangling of these components by quantifying the representation quality at each layer, independently of the specific unembed module being used for vocabulary readout.

The prompting strategies compared in the work include instruction prompts and demonstration prompts. As a baseline control for the representation analysis, the authors also extracted embeddings using the raw sentence input. The paper considers sentence embeddings (mean of residual stream activations for tokens corresponding only to the input sentence) and last-token embeddings (residual stream activations of the last token in the sequence at each layer) to analyze representational geometry in decoder-only models. Category manifolds are constructed by accumulating the embedding vectors of all sentences sharing a class label. Manifold capacity is measured as a scalar, measuring how separable the underlying category manifolds are. The dimension of each manifold was measured as the participation ratio of principal components, while the radius of each manifold was taken to be the maximum distance between any pair of points on the manifold. Correlation coefficients between axes of variation of individual manifolds (Axes-alignment) and correlations between each manifold's axes and its centroid (Center-axes alignment) are measured as the correlation structure to explain capacity changes driven by the relative arrangements of manifolds in the embedding space.

The analysis reveals complex dynamics in how prompting affects the internal representations of LLMs. Instruction alone achieved good accuracy, outperforming demonstration prompts with few examples. Larger example sets surpassed explicit instruction, with performance quickly plateauing. Replacing meaningful category words with abstract letters required more demonstration examples to infer category nature. When category labels were consistently shuffled, the model failed to generalize beyond pretrained associations, achieving low accuracy for both target and original labels. Analysis of sentence-level embeddings revealed that demonstration examples, but not abstract instruction, significantly reorganized intermediate representations at early-mid layers. This reorganization increased the separability of sentiment manifolds by reducing manifold dimension and improving correlation structure. There was little difference in resulting geometry between demonstrations across three labeling strategies, indicating that sentence representation is primarily influenced by input distribution examples, rather than input-output mapping. At the last-token level, instruction prompts significantly increased manifold capacity relative to raw sentences, with effects emerging as early as layer 8 and persisting to final layers. Demonstrations further increased manifold capacity compared to instruction, despite lower task performance for cases with few examples. Last-token capacity during letter code labeling was much lower compared to category words, explaining lower performance when output labels lack meaningful semantics. For shuffled labels, capacity values were similar to the gold label setting, suggesting that the model's inability to overwrite existing associations is explained by the readout misalignment, while the underlying representation is intact. Despite substantial performance variability due to the choice of particular examples and their ordering, the changes in manifold capacity of the last token embedding at the final layer were minimal.

In a multi-task setting, the authors investigated whether prompting a model to perform one task would affect the quality of representation for another unrelated task. Increasing the number of demonstrations robustly led to increased manifold capacity at intermediate layers for coherent configurations, while instruction had a much weaker effect. Demonstrations for an incoherent task also increased capacity with a similar layerwise profile, highlighting the role of input distribution. Analysis of last token embeddings revealed that at earlier layers, additional demonstrations of incoherent tasks increased manifold capacity, but at later layers, this trend reversed, with additional examples decreasing capacity. This suggests a representational tradeoff: as the model prepares the output, features for irrelevant tasks, that were emphasized at intermediate processing stages are compromised in favor of better separability of task-relevant features.

The investigation was extended to prompt-tuning, which optimizes a task-specific prompt directly in the embedding space, prepended to the test input. Soft prompts consistently outperformed demonstrations, and there was no significant correlation between final performance and prompt length. The optimization-based solution did not alter intermediate representations in earlier layers. At the last token level, soft prompts dramatically reduced the capacity of representations for unrelated tasks in the incoherent case. Soft-prompts operate through fundamentally different internal mechanisms compared to demonstrations and zero-shot instruction.

Zero-shot instructions primarily influence the final stages of processing, affecting how features are "packaged" in the last token embedding without significantly altering intermediate representations. Demonstration examples have a more profound impact, reshaping intermediate representations to optimize them for the classification objective. Soft-prompts mainly affect later layers responsible for output preparation, distinguishing them from the broader impact of natural language demonstrations. A key insight is the distinction between representational geometry and readout alignment in determining model performance. Internal representations often maintain high manifold capacity even when ICL (In-Context Learning) performance is poor. There is significant potential in better understanding and optimizing readout alignment. The success of prompt-tuning stems primarily from improving the alignment between representations and the vocabulary readout layer, rather than fundamentally altering the geometric organization of the embedding space. The authors suggest that internal representations often maintain high manifold capacity even when ICL (In-Context Learning) performance is poor, indicating potential for better understanding and optimizing readout alignment.

The paper used synthetic datasets generated by Claude 3.5 Sonnet, which may not fully capture the complexity and variability of real-world language structure. The metrics used to quantify representational geometry, such as manifold capacity and individual manifold geometry, simplify the more complex tasks that occur in LLMs. Future work should examine how other tasks, such as those requiring multi-token outputs, affect representational geometry.