Skip to main content

A plug-and-play, scikit-learn compatible implementation of Supervised Multi-Dimensional Scaling (SMDS) for automatic feature manifold discovery in LLMs.

Project description

Supervised Multi-Dimensional Scaling

Arxiv License Python Versions CI

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}, 
}

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

smds-1.6.5.tar.gz (4.4 MB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

smds-1.6.5-py3-none-any.whl (73.6 kB view details)

Uploaded Python 3

File details

Details for the file smds-1.6.5.tar.gz.

File metadata

  • Download URL: smds-1.6.5.tar.gz
  • Upload date:
  • Size: 4.4 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.10.7 {"installer":{"name":"uv","version":"0.10.7","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for smds-1.6.5.tar.gz
Algorithm Hash digest
SHA256 1cf22bacfcb12d2518c4e001dc5b0cc7eddef9067ce9f73f6ae9a69dff33ccaa
MD5 c0774d7150dffa66afbf72ae6380f6dc
BLAKE2b-256 f277d5a24eea7d242e137d5564ca7bba48e0796454ed793de6d08d255062d5a9

See more details on using hashes here.

File details

Details for the file smds-1.6.5-py3-none-any.whl.

File metadata

  • Download URL: smds-1.6.5-py3-none-any.whl
  • Upload date:
  • Size: 73.6 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.10.7 {"installer":{"name":"uv","version":"0.10.7","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for smds-1.6.5-py3-none-any.whl
Algorithm Hash digest
SHA256 62a0ad622b7ca208e2dfcb01ce501572f5d6a10b0231fb8c615afe9ca5f09825
MD5 47d93b93a7dc63040cde77630510e270
BLAKE2b-256 7c33b9c9d93b840af5fee1cc5d905af614899d05347266182a4abdca5a7f1fc6

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page