Supervised Multi-Dimensional Scaling

This is a stand-alone implementation of Supervised Multi-Dimensional Scaling (SMDS) from the paper "Shape Happens: Automatic Feature Manifold Discovery in LLMs". It contains a plug-and-play class written with the familiar scikit-learn interface. SMDS supports several template shapes to discover manifolds of various shape.

Contact person: Federico Tiblias

UKP Lab | TU Darmstadt

Don't hesitate to report an issue if you have further questions or spot a bug.

Getting started

With uv (recommended):

uv add smds

With pip:

pip install smds

Usage

The SupervisedMDS class provides a scikit-learn style interface that is straightforward to use. Unlike standard MDS, it requires selecting a stage-1 strategy and target manifold by name (for example: "cluster", "circular").

Fit & Transform

You can instantiate the model, fit it to data (X, y), and transform your input into a low-dimensional embedding:

import numpy as np
from smds import SupervisedMDS

# Example data
X = np.random.randn(100, 20)  # 100 samples, 20 features
y = np.random.randint(0, 5, size=100)  # Discrete labels (clusters)

# Instantiate and fit
# stage_1: "computed" (default) or "user_provided"
# manifold: one of the built-in shape names, e.g. "cluster", "circular", "log_linear"
smds = SupervisedMDS(stage_1="computed", manifold="cluster", alpha=0.1)
smds.fit(X, y)

# Transform to low-dimensional space
X_proj = smds.transform(X)
print(X_proj.shape)  # (100, 2)

If you set stage_1="user_provided", manifold is ignored and a warning is raised.

Manifold Discovery

Once fitted, you can use the learned transformation for inverse projections and to assess how well the embedding matches the target geometry:

from smds.pipeline.discovery_pipeline import discover_manifolds
from smds.pipeline import open_dashboard

# Run discovery pipeline
# Evaluates default shapes (Cluster, Circular, Hierarchical, etc.)
# Returns a DataFrame sorted by best fit (lowest stress / highest score)
df_results, save_path = discover_manifolds(
    X,
    y,
    smds_components=2,  # Target dimensionality
    n_folds=5,  # Cross-validation folds
    experiment_name="My_Exp",  # Name for saved results
    n_jobs=-1  # Use all available cores
)

print(f"Best matching shape: {df_results.iloc[0]['shape']}")
print(df_results.head())

# Launch the interactive Streamlit dashboard to explore results and plots
open_dashboard.main(save_path)

The discovery pipeline handles:

• Hypothesis Testing: Iterates through a default or custom list of manifold shapes.
• Cross-Validation: Uses k-fold CV to ensure robust scoring.
• Caching: Caches intermediate results to resume interrupted experiments.
• Visualization: Generates interactive plots for the dashboard.

Statistical Validation

Standard cross-validation provides a mean score, but it does not tell you if one manifold is statistically better than another. SMDS includes a robust Statistical Testing (ST) wrapper that runs repeated experiments to perform a * Friedman Rank Sum Test* and Nemenyi Post-Hoc Analysis.

Running a Statistical Test

Instead of smds/pipeline/run_pipeline, use the smds/pipeline/run_statistical_test.py wrapper:

from smds.pipeline.statistical_testing.run_statistical_test import run_statistical_validation

# Runs the pipeline 10 times (10 repeats), each with 5-Fold CV
pivot_dfs, output_path = run_statistical_validation(
    X=my_data,
    y=my_labels,
    n_repeats=10,
    n_folds=5,
    experiment_name="my_robust_experiment"
)

Viewing Results

Open the dashboard to view the Friedman Statistic, P-Value Heatmap, and Critical Difference (CD) Diagram:

from smds.pipeline import open_dashboard

# Launch the interactive Streamlit dashboard to explore results and plots
open_dashboard.main()

Testing

Run the test suite using pytest:

  make test

Optimization & GPU Support

For manifolds with undefined distances (e.g. ChainShape), SMDS falls back to a generic SciPy solver.
For large datasets, this can be slow.

SMDS provides an optional accelerated solver based on PyTorch, which is significantly faster on CPU and can transparently leverage GPUs when available.

Enabling the Accelerator

Install SMDS with the optional fast extra:

  pip install smds[fast]

Then enable it in your model:

smds = SupervisedMDS(
    ...,
    manifold="chain",
    gpu_accel=True,
)

If a compatible GPU is available, PyTorch will use it automatically. Otherwise, the accelerated solver will run on CPU.

💡 GPU support (CUDA on NVIDIA, MPS on Apple Silicon) depends on your PyTorch installation. See the official PyTorch documentation for platform-specific setup.

Documentation

The full online documentation is available at SMDS Documentation.

To build and serve the documentation locally:

mkdocs serve

Contributors (the Shape Wizards team)

Contribution Matrix

Category	Anton	Arwin	Jan	Simon	Vinayak
Shape Implementation	Geodesic, Cylindrical, Spherical	KleinBottle, HierarchicalShape	Shapes: CircularShape SpiralShape	ClusterShape, DiscreteCircular, ChainShape	LogLinear, Euclidean, SemiCircular, Torus, Polytope, GraphGeodesic
Architectural Extensions	BaseShape API design/implementation, Precomputed manifolds (Y bypass)	BaseShape, AlternativeReducer	Design base shape abstraction. Design of stress metrics following sklearn conventions following research papers
Stress Metrics		Non-metric stress	Implement abstraction of stress metric. Shepard score.		Normalized stress, KL divergence
Discovery, Visualization & Validation		Discovery Pipeline, png-Visualization		Discovery Pipeline, Dashboard, Statistical Testing
Infrastructure, Tooling, Testing & Performance	Default normalization in BaseShape, General refactoring (PR) Shapes API	Unified stress tests, sklearn compatibility tests	Leading engineering efforts: Implementing CI/CD pipelines Implementing environment reproducibility Selection of tools used by our team for coding and organisation Deploying package to PyPi Designing rules for Pull Requests review process Creating documentation	Shape Integration Test Framework, GPU acceleration for ChainShape
Experiments	Manifold Activation Patching	Hour of Manifold Experiment	Experimental notebook	Manifold Location Experiment	Color Manifolds in VLMs

Cite

Please use the following citation:

@misc{tiblias2025shapehappensautomaticfeature,
      title={Shape Happens: Automatic Feature Manifold Discovery in LLMs via Supervised Multi-Dimensional Scaling}, 
      author={Federico Tiblias and Irina Bigoulaeva and Jingcheng Niu and Simone Balloccu and Iryna Gurevych},
      year={2025},
      eprint={2510.01025},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2510.01025}, 
}