torch-supercluster

SuperCluster learns K prediction prototypes jointly with a supervised objective. Each example is routed to prototypes via cross-attention, and the prediction is an exact convex combination of prototype-level predictions.

This gives you a model that is simultaneously:

• Accurate — matches or exceeds MLP accuracy on tabular benchmarks
• Interpretable — each prototype maps to a readable "prediction archetype"
• Self-analysing — over-specify K and the model tells you how many distinct prediction regimes the data contains

x → MLP encoder → cross-attention (→ prototype centers C) → soft routing π
prediction:  ŷ = σ(π · s)   [binary]
             softmax(π @ S)  [multi-class]
             π · s           [regression]

Installation

pip install torch-supercluster

Requires Python ≥ 3.9 and PyTorch ≥ 2.0.

Quick start

Binary classification

from supercluster import SuperCluster, center_diversity_loss
import torch, torch.nn as nn

model = SuperCluster(
    input_dim=10,
    embedding_dim=64,
    num_clusters=8,       # set larger than expected; surplus prototypes collapse
)

logits, cluster_weights = model(x)               # x: [N, 10]
loss = nn.BCEWithLogitsLoss()(logits, y)          # logits: [N, 1]
loss = loss + 0.1 * center_diversity_loss(model.centers)
loss.backward()

Multi-class classification

model = SuperCluster(
    input_dim=54, embedding_dim=128,
    num_clusters=8, num_classes=7,
)
logits, cluster_weights = model(x)               # logits: [N, 7]
loss = nn.CrossEntropyLoss()(logits, y_long)

Regression

model = SuperCluster(
    input_dim=10, embedding_dim=64, num_clusters=5,
)
preds, cluster_weights = model(x)                # preds: [N, 1], no sigmoid
loss = nn.MSELoss()(preds, y)

Reading prototype predictions

import torch.nn.functional as F

# Binary: prototype predicted probability
for k in range(model.num_clusters):
    prob = torch.sigmoid(model.prototype_scores[k]).item()
    print(f"Prototype {k}: P(y=1) = {prob:.3f}")

# Multi-class: prototype predicted class distribution
for k in range(model.num_clusters):
    probs = F.softmax(model.prototype_scores[k], dim=0)
    print(f"Prototype {k}: {probs.tolist()}")

How it works

SuperCluster separates two concerns that prior prototype models conflate:

The center-diversity loss pushes centers apart on the unit hypersphere without touching prediction scores, preventing the "overconfident prototype" pathology common when routing geometry and prediction are conflated.

Predictive regime discovery

Set K larger than you expect and examine which prototypes are occupied. The model concentrates assignment mass on exactly as many prototypes as the data has distinct prediction behaviours — surplus prototypes collapse to zero routing mass. This gives an operational measure of how many prediction regimes the data contains.

from supercluster import effective_prototype_count

_, cw = model(X_test)
k_eff = effective_prototype_count(cw)          # entropy-based count
active = model.active_prototypes(cw)           # list of occupied prototype indices
print(f"K_eff = {k_eff:.2f}, K_active = {len(active)}/{model.num_clusters}")

Model parameters

Training recommendations

Empirical results

Citation

If you use this package in academic work, please cite:

@article{danielson2026supercluster,
  title   = {SuperCluster: Learning Prediction Prototypes via
             Target-Guided Cross-Attention Clustering},
  author  = {Danielson, Aaron John},
  journal = {Data Mining and Knowledge Discovery},
  year    = {2026},
}

License

MIT