TL;DR
This paper introduces Supervised Multi-Dimensional Scaling (SMDS), a novel method to analyze how language models (LMs) represent concepts as geometric structures (manifolds) in their internal workings. They demonstrate that these manifolds are not random but reflect semantic properties, actively support reasoning, and adapt to context, offering a new lens into how LMs understand and process information.
What Problem Does This Paper Address?
The "linear representation hypothesis" suggests that LMs encode concepts as directions in their latent space, forming organized "feature manifolds." While prior work has identified some of these structures, it often focused on individual features or assumed specific geometries, limiting generalization. A key challenge has been to systematically evaluate and compare different hypothesized manifold structures across various LMs and tasks. Furthermore, understanding how LMs utilize these manifolds for reasoning and how they adapt to different contexts remains an open question. This paper aims to bridge this gap by providing a flexible and quantitative method to analyze these feature manifolds.
Methods
The core contribution is Supervised Multi-Dimensional Scaling (SMDS). It's a model-agnostic dimensionality reduction technique that extends classical Multi-Dimensional Scaling (MDS). Unlike standard MDS, SMDS incorporates supervision by using labels to define an ideal geometric structure for a feature manifold. It then learns a linear projection from the LM's high-dimensional activations to this hypothesized low-dimensional manifold. The method reframes manifold identification as a model selection problem, allowing researchers to quantitatively evaluate how well different geometric hypotheses (e.g., linear, circular, clustered) fit the data. This allows for a systematic comparison of competing manifold structures and provides a goodness-of-fit metric.
Datasets
The authors created synthetic datasets specifically for temporal reasoning tasks. These datasets cover various temporal expression types:
- date: Calendar dates.
- date_season: Dates categorized by season.
- date_temperature: Dates categorized by temperature (implied by season).
- duration: Length of time.
- notable: Indirect temporal references to specific events.
- time_of_day: Specific times within a day.
- time_of_day_phase: Time of day categorized into phases (e.g., morning, evening).
They also mention using a spatial reasoning dataset involving city locations.
Key Insights and Ideas
- Manifolds are semantically meaningful: The geometric shapes of feature manifolds (circles for dates, lines for durations) directly correspond to the semantic properties of the concepts they represent.
- Universality across LMs: These manifold structures are remarkably consistent across different LM families and sizes, suggesting a common underlying organizational principle.
- Dynamic adaptation: LMs don't have static representations; feature manifolds actively reshape in response to the specific task and contextual cues provided in the prompt.
- Crucial for reasoning: Perturbing these identified manifold subspaces significantly degrades LM performance on reasoning tasks, while perturbing random subspaces has little effect. This provides causal evidence that these manifolds are actively used for reasoning.
- Information flow: Manifold structures are not confined to specific tokens but are preserved and propagated throughout the LM's layers, indicating a mechanism for binding information.
Key Results
- SMDS successfully identified interpretable geometric structures (circles, lines, clusters) for temporal entities across various LMs (Qwen2.5, Gemma-2, Llama-3 variants).
- The best-fitting manifold structures were consistent across different model sizes (3B, 8B, 70B) and architectures.
- Experiments showed that LMs dynamically adjust feature manifolds based on task context, for instance, using circular structures for dates but linear or clustered structures for seasonal or temperature-based reasoning.
- Causal interventions demonstrated that disrupting manifold-aligned subspaces significantly impaired downstream accuracy, while random subspace perturbations had minimal impact. For example, noise injection into a low-dimensional manifold subspace was as effective at degrading performance as perturbing much larger activation spaces.
- Manifold quality (measured by stress) showed a significant positive correlation with downstream accuracy, particularly for models that already achieved above-chance performance.
Limitations and Future Work
- Bias in corpora: The use of English corpora and the Gregorian calendar might introduce biases in the learned representations.
- Fuzzy expressions: The study omitted fuzzy temporal expressions (e.g., "later," "next week") due to the difficulty in defining precise temporal pointers, which could impact performance degradation.
- Tightly clustered scores: SMDS stress scores can sometimes be tightly clustered, making it challenging to definitively identify a single best manifold. The authors addressed this with bootstrapping and statistical testing, but acknowledge that models might adopt multiple valid geometries.
- Future directions:
- Investigating how individual features combine into multidimensional manifolds.
- Developing more discriminative metrics for manifold analysis.
- Exploring the coexistence of multiple representational geometries and how LMs select them.
- Combining SMDS with circuit discovery to understand how LMs transform information.
- Using manifold knowledge for model steering.
- Further studying the impact of noise on feature manifolds and its relation to reasoning.
- Extending the analysis to fuzzy temporal expressions.