Shape Happens: Feature Manifold Discovery

TL;DR

Language models don’t just encode concepts as individual directions — they organize related features into structured geometric shapes (circles, chains, clusters). This paper introduces SMDS (Supervised Multi-Dimensional Scaling), a framework that automatically discovers and compares these geometric hypotheses from labeled activation data, and shows the shapes are semantically meaningful and functionally active during inference.

Abstract

The linear representation hypothesis posits that language models encode concepts as directions in their latent spaces, forming organized multidimensional manifolds. We introduce Supervised Multi-Dimensional Scaling (SMDS), a model-agnostic technique for testing and comparing competing hypotheses about how LLMs organize concepts. We show that feature manifolds display diverse geometric patterns — circles, lines, clusters — that reflect the semantic properties of the represented concepts, remain consistent across model families and sizes, and actively participate in model reasoning, adapting dynamically when contextual information changes.

Reproducing with Murano

The original codebase runs the full analysis through a series of standalone scripts: extracting activations, running probing experiments, fitting SMDS, and plotting results. With Murano, the activation recording and probing steps collapse into a single pipeline.

Step 1 — Record activations and find the best layer

from murano import LabeledDataset, MuranoModel, Pipeline
from murano.steps import Load, Record, Probe

model = MuranoModel("meta-llama/Llama-3.2-1B-Instruct")

# Load a labeled dataset (e.g. month-of-year concept)
dataset = LabeledDataset.from_hub(
    "coastalcph/fem-str-0-en",
    text_column="text",
    label_column="label",
    n=1000,
)

results = Pipeline([
    Load(dataset),
    Record(model, layers="all", position="last"),
    Probe(cv=5),
]).run()

probe   = results["probe"]
store   = results["record"]
best    = probe.best_layer

print(f"Best layer: {best}  |  Accuracy: {probe.accuracy_per_layer[best]:.3f}")

Murano handles tokenization, batching, layer indexing, and cross-validation. The equivalent in the original codebase spans three scripts and requires manual alignment of activation tensors with label arrays.

Step 2 — Discover the feature manifold geometry

from smds import SupervisedMDS
import numpy as np

X = store.activations[best]          # (n_samples, d_model)
y = np.array(store.labels)

smds = SupervisedMDS(stage_1="computed", manifold="auto")
smds.fit(X, y)
X_proj = smds.transform(X)           # (n_samples, 2)

print(f"Best-fit manifold: {smds.best_manifold_}")

manifold="auto" runs cross-validated selection across all candidate shapes (cluster, circle, chain, hierarchy, …) and returns the best fit via a Friedman rank test. The paper finds a circular manifold for temporal concepts like months or days, and cluster structure for categorical concepts like sentiment.

Step 3 — Visualize

import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(5, 5))
scatter = ax.scatter(X_proj[:, 0], X_proj[:, 1], c=y, cmap="tab10", s=18)
ax.set_title(f"SMDS projection — {smds.best_manifold_} manifold")
plt.colorbar(scatter, ax=ax, label="Label")
plt.tight_layout()
plt.show()

Full pipeline (combined)

from murano import LabeledDataset, MuranoModel, Pipeline
from murano.steps import Load, Record, Probe
from smds import SupervisedMDS
import numpy as np, matplotlib.pyplot as plt

model   = MuranoModel("meta-llama/Llama-3.2-1B-Instruct")
dataset = LabeledDataset.from_hub(
    "coastalcph/fem-str-0-en",
    text_column="text",
    label_column="label",
    n=1000,
)

results = Pipeline([Load(dataset), Record(model, layers="all", position="last"), Probe(cv=5)]).run()
probe, store = results["probe"], results["record"]
best = probe.best_layer

X = store.activations[best]
y = np.array(store.labels)

smds = SupervisedMDS(stage_1="computed", manifold="auto")
smds.fit(X, y)
X_proj = smds.transform(X)

fig, ax = plt.subplots(figsize=(5, 5))
ax.scatter(X_proj[:, 0], X_proj[:, 1], c=y, cmap="tab10", s=18)
ax.set_title(f"SMDS — {smds.best_manifold_} manifold (layer {best})")
plt.tight_layout(); plt.show()

The dataset slug coastalcph/fem-str-0-en is a placeholder — replace it with the exact HuggingFace dataset used in the paper. The Colab notebook above runs the full experiment end-to-end with the correct datasets.

Key results

Concept type	Model	Best-fit manifold	Probe accuracy
Month of year	Llama-3.2-1B	Circle	0.91
Day of week	Llama-3.2-1B	Circle	0.87
Sentiment	Llama-3.2-1B	Cluster	0.94
Country → continent	Llama-3.2-1B	Hierarchy	0.88

Manifold geometry is consistent across Llama, Mistral, and Gemma families and scales with model size.

Citation

@article{tiblias2026shape,
  title     = {Shape Happens: Automatic Feature Manifold Discovery
               in LLMs via Supervised Multi-Dimensional Scaling},
  author    = {Tiblias, Federico and Bigoulaeva, Irina
               and Niu, Jingcheng and Balloccu, Simone
               and Gurevych, Iryna},
  journal   = {Transactions on Machine Learning Research},
  year      = {2026},
  url       = {https://arxiv.org/abs/2510.01025}
}

Shape Happens: Automatic Feature Manifold Discovery in LLMs via Supervised Multi-Dimensional Scaling