Skip to main content

PyTorch and Affinity Propagation

PyTorch 2.14.1 does not provide a maintained high-level Affinity Propagation estimator. Handwritten message loops are easy to get wrong: the responsibility maximum must exclude the candidate, availability has separate diagonal and off-diagonal rules, and damping must retain old messages. Use PyTorch for tensor preparation and a tested estimator for clustering.

1. Data and tensor similarity​

import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.cluster import AffinityPropagation
from sklearn.datasets import make_blobs
from sklearn.metrics import adjusted_rand_score

X, reference = make_blobs(n_samples=180, centers=3, cluster_std=0.60, random_state=42)
X_tensor = torch.tensor(X, dtype=torch.float64)
squared_distances = torch.cdist(X_tensor, X_tensor).square()
similarity_tensor = -squared_distances
assert torch.allclose(similarity_tensor, similarity_tensor.T)
similarity = similarity_tensor.numpy()

2. Choose an explicit preference and fit​

off_diagonal = similarity[~np.eye(len(X), dtype=bool)]
preference = float(np.median(off_diagonal))
model = AffinityPropagation(
    affinity="precomputed",
    preference=preference,
    damping=0.8,
    random_state=42,
).fit(similarity)
labels = model.labels_
assert model.n_iter_ < model.max_iter
assert adjusted_rand_score(reference, labels) > 0.95

3. Plot observations and exemplars​

exemplars = X[model.cluster_centers_indices_]
plt.scatter(X[:, 0], X[:, 1], c=labels, cmap="viridis", s=24)
plt.scatter(exemplars[:, 0], exemplars[:, 1], c="red", marker="x", s=100)
plt.title("Affinity Propagation with a PyTorch similarity")
plt.close()

The CPU conversion is explicit. A GPU can accelerate distance construction, but the dense matrix still needs quadratic memory and scikit-learn fits on CPU.