What it actually does
Imagine every data point is a voter and every data point is also a candidate. Each voter tells each candidate how much it would like that candidate to represent it, measured against its best alternative — that is the responsibility. Each candidate answers with how available it is, which depends on how much support it is already collecting from everyone else — that is the availability. The two messages are passed back and forth until the votes settle, and the points that end up voting for themselves are the exemplars. Nobody told the algorithm how many clusters to find; that falls out of how reluctant each point is to stand as a candidate, a number called the preference. The leaders are real data points, not averages, which means the only thing the method ever needs is a table of how similar each pair is — and that table can come from edit distance, from correlation, from anything that can be computed between two items.
- A pairwise similarity exists and is meaningful for the data; it need not be a metric, symmetric, or computed from vectors, but it is the whole model and a poor one cannot be repaired downstream
- Each cluster has a point that represents it well — an exemplar that is an actual member, which fails for ring-shaped or elongated clusters where no single member is central
- A single preference (or a sensible per-point vector of them) suits the whole dataset, so clusters of very different size and density are not all expressible at once
- The data is small enough for an n-by-n table, or a sparse neighbour similarity is a faithful substitute for the dense one
- The message passing converges under the chosen damping; this is an assumption about the data and the settings, not a property of the algorithm
Fixed-point condition
Choose an exemplar for every point so that the total similarity between points and their exemplars is as large as possible, with one hard rule: a point may only be chosen as an exemplar if it chooses itself. Because each exemplar scores its own preference, every extra cluster costs something, and that cost is what stops the answer being “everyone is their own exemplar”. That is the target, and it is a combinatorial problem of the facility-location family, NP-hard in general. The method does not descend it. It is derived by writing the problem as a factor graph and running max-sum belief propagation, and what it actually computes is a fixed point of the two message updates, from which the exemplars are read off. That is why this entry says fixed-point rather than loss: nothing monotone is being minimized — net similarity can go down between iterations — and the algorithm’s output is whatever the messages settle to. Where the messages do settle, the result is empirically near-optimal; the general theory of max-product on loopy graphs gives optimality only over a restricted neighbourhood of assignments, which is weaker than any approximation ratio, so no guarantee against the true optimum is available or should be claimed.
- the exemplar that point i is assigned to; c_i = i means i is itself an exemplar
- how well the chosen exemplar serves point i — a similarity, so larger is better
- the preference: what an exemplar “scores” for itself, which makes it the price of opening a cluster and the only control over how many appear
- the exemplar-consistency penalty: infinite if anyone chooses k while k does not choose itself, so a point cannot be followed unless it leads
How it is minimized
Max-sum message passing on the exemplar factor graph, with damped fixed-point iteration
- responsibility sent from i to k: how strongly i wants k as its exemplar, relative to the best rival i has
- availability sent from k to i: how appropriate it would be for i to choose k, given the support k already has
- self-responsibility: evidence that k should be an exemplar at all, seeded by the preference s(k,k)
- damping factor in [0.5, 1): the fraction of the old message kept, applied to both r and a — the thing that stops the iteration oscillating
- the decoded assignment; k is an exemplar when a(k,k) + r(k,k) is positive
Written as belief propagation, the clustering problem becomes two local rules, which is why the method exists. The responsibility rule is a competition: i compares k against its best rival, so a point is drawn toward the candidate that most clearly beats the alternatives. The availability rule is accumulation of evidence: k becomes more available to i the more positive responsibility other points direct at it, capped at zero from above so that support is only ever withdrawn, never invented. The three things that carry real consequences are not in the formulas. The cost is quadratic per iteration only because of an implementation trick — keep the largest and second-largest value in each row, and the max over rivals is a lookup, whereas transcribing the formula costs O(n) per entry and O(n³) per sweep. Damping is not optional polish: undamped max-sum on this graph oscillates on many real inputs, and the symptom looks like a data problem rather than a numerical one. And exact ties between similarities leave the messages with no way to break symmetry, so a tiny relative jitter is added — the reference implementations all do this, and it is a correctness measure rather than a tweak.
similarity function s(i, k)
The model. Default is negative squared Euclidean distance, which turns the objective into a facility-location version of squared error. Any similarity is allowed — negative edit distance for sequences, cosine for documents, Pearson correlation for series, even an asymmetric directed cost — because nothing requires a metric, a mean, or symmetry. The sign matters: passing a distance where a similarity is expected selects the most outlying points as exemplars (any real-valued table; larger means more alike, and the scale sets the scale of the preference)
preference s(k, k)
The price of an exemplar, and the only control over how many clusters emerge. A higher value yields more clusters, a lower one fewer. It lives in the units of the similarity, so changing the similarity changes what every value means, and it may be a vector, giving individual points a head start (median of the similarities for a moderate count, minimum for few clusters; scan on a log-like scale)
damping λ
How much of the previous message survives each update. Too low and the messages oscillate without converging; too high and convergence is slow but steady (0.5 default, commonly raised to 0.7 to 0.95 when oscillation appears)
max iterations
The hard budget. Reaching it means the messages did not settle, which is a result to act on rather than to accept (200 to 1000)
convergence window
Stop when the set of exemplars has not changed for this many consecutive iterations. The message values keep drifting long after the decision is stable, so the decision is what is monitored (15 to 50 iterations)
There is no convergence guarantee. Max-sum belief propagation on a graph with loops can oscillate indefinitely, and it does so on real data; damping suppresses it and is the first remedy, with a higher λ the second, and a few percent of jitter in the similarities the third. There is also no monotone quantity to watch, because the objective is not what is being descended — the reliable signal is the exemplar set stabilizing. The characteristic failures are all visible once you know to look. Exemplars that flip between iterations mean the damping is too low. A result with one giant cluster or with nearly every point an exemplar means the preference is on the wrong side of the similarity scale. No exemplars at all means the preference is so low that no point believes in itself. And a similarity matrix with a sign error returns confident, tidy, entirely wrong exemplars. At a fixed point the answer is typically very good, and it should still be judged against a baseline such as k-medoids rather than trusted to be optimal.
Per iteration: O(n²) time with the top-two trick, O(n³) if the max over rivals is transcribed literally. Memory: three n-by-n matrices — similarity, responsibility, availability — plus temporaries, which is the real ceiling: roughly ten thousand points on a workstation, with the matrices a gigabyte or so each in double precision. Iterations to converge are typically in the hundreds. On a sparse similarity — each point connected only to its nearest neighbours, every absent entry treated as minus infinity — both time and memory drop to O(edges) per iteration, which is how the method is pushed well past that ceiling, at the cost of whatever the neighbour pruning discards.
How, where, and why it is used
Time-Series Forecasting
AdaptedCluster many series by how they move, and pick a real series from each group as its representative. Build the similarity as the correlation between series after differencing and z-normalizing, or as negative dynamic-time-warping distance; neither admits a mean, which is the reason this suits series better than k-means. The exemplar is an actual series, so it can be plotted, inspected, forecast directly, or used as the anchor for a pooled model across its group. An asymmetric variant is possible: let s(i, k) be the best lagged correlation of i with k leading it, so the table itself says which series are leaders and which are followers.
- Grouping thousands of related series — store-by-product demand, per-asset metrics — so one model is fitted per group instead of one per series
- Picking representative series for a dashboard or a review, where an average curve would be misleading and a real one is not
- Lead-lag structure over sensors or markets using an asymmetric lagged-correlation similarity
It accepts exactly the kind of similarity series produce — correlation or time-warped distance, with no meaningful average — and it returns real series as cluster centres, which makes the output explainable. It also lets a new series be assigned cheaply: take the exemplar it is most similar to, with no refit. The costs bind quickly: an n-by-n correlation table over tens of thousands of series is not affordable, results shift when the series set changes because the exemplar set can change, and the number of groups is controlled by a preference that must be tuned on a scale defined by the correlations rather than by a count anyone chose.
- Difference or detrend before computing correlation: two independently trending series correlate strongly for no reason, and the clusters recover trend rather than behaviour
- Z-normalize each series, or the similarity measures level and scale and the clusters group by magnitude
- Set the preference from the correlation table (its median, or a quantile) rather than from a number carried over from another dataset
- Compute the similarity on the training window only, or the grouping leaks future co-movement into the forecast
Rolling-origin backtest of the forecasts made with the clustering in place — pooled-per-group versus per-series — not any clustering score. Check that the number of groups is stable across refits and that exemplar identities do not churn.
- Spurious correlation between trending series, which makes the table describe shared trend rather than shared dynamics
- A preference that over-fragments, leaving most groups with one or two series and nothing to pool
- Exemplar churn across refits, which breaks any downstream pipeline keyed on the exemplar’s identity
- Quadratic memory reached silently as the series count grows between refits
Anomaly Detection
Not applicableThe method has no notion of outlyingness, and its own mechanism works against it: a point dissimilar to everything has no better option than to vote for itself, so an outlier typically becomes a singleton exemplar — a tidy cluster that hides it — whenever its similarities fall below the preference. The only score available is similarity to the assigned exemplar, which is the generic distance-to-centre recipe any clustering offers and not something this method contributes. Use an isolation forest, a nearest-neighbour score, or DBSCAN’s noise label instead.
Optimization
AdaptedThe stated objective is the uncapacitated facility-location problem in disguise: choose which points to open as facilities (exemplars), assign every point to one open facility, maximize total similarity — equivalently minimize total assignment cost — where the preference is the cost of opening a facility. Set s(i, k) to the negative cost of serving i from k, s(k, k) to the negative opening cost, and the number of facilities is chosen by the economics rather than fixed in advance. Asymmetric costs are natural here, since the cost of serving i from k need not equal the reverse. The algorithm is then a heuristic for that combinatorial problem, with the structure of its output — real points as centres — matching the structure of the decision.
- Depot, warehouse or hub placement from a table of travel costs, which is asymmetric when routes are one-way or congested
- Choosing a representative subset — prototypes for a labelling budget, or a reduced set of scenarios in stochastic planning — where the selected items must be real members of the set
- Exemplar and prototype selection for summarization and for seeding other methods with real starting points
Compared with k-medoids it needs no number of facilities in advance and does not depend on a random start; compared with an exact solver it scales more gracefully on a similarity table that is already given. Against it, and decisively when the instance is small enough: facility location has exact mixed-integer formulations and LP-rounding algorithms with approximation guarantees, and this method has neither a bound nor a way to add real constraints such as capacities. If a solver can handle the instance, use the solver; this earns its place when the cost table is large, soft, and the answer need only be good.
- Encode costs as negative similarities in consistent units, and set the preference to the real opening cost in those same units — they are the same quantity
- Encode infeasible assignments as minus infinity entries, which both forbids them and, on a sparse table, saves the memory
- Capacities and per-site limits cannot be written into the similarity; check them afterwards and treat a violation as a reason to switch method
Total cost of the returned solution against a solver or an LP-relaxation lower bound where the instance allows it, and against k-medoids and a greedy baseline where it does not. Report the gap, not just the cost.
- A preference in the wrong units, which silently misprices the number of facilities by an order of magnitude
- Assuming the answer is optimal because it converged; a fixed point is not a certificate
- Treating unmodelled constraints as someone else’s problem, then discovering the returned sites cannot be built
Natural Language
ViableCluster documents, sentences or short texts using cosine similarity over tf-idf vectors or sentence embeddings, and report each cluster by its exemplar — which is a real piece of text, not a centroid vector that must be decoded back into words. Used on sentences from a document, the exemplars are an extractive summary.
- Extractive summarization of a document or a thread, taking exemplar sentences as the summary
- Grouping support tickets or news items, labelling each group with its most representative real item
- Near-duplicate collapse in a text corpus, where the cluster’s exemplar stands in for the group
Cosine is the right similarity because it compares direction and ignores length, so a short and a long document on one topic are close; swap in Euclidean distance over raw counts and clusters form by document length instead of topic. Exemplars that are readable text are the practical advantage over k-means, whose centroid is a vector no one can read. The costs are the usual quadratic ones, and a specific trap: embedding similarities tend to occupy a narrow band near the top of the scale, so the preference has little room and a small change swings the cluster count.
- L2-normalize the vectors so a dot product is the cosine, and compute the full similarity in one matrix product
- Remove boilerplate and duplicates first, or the exemplars are the boilerplate
- Inspect the distribution of similarities before choosing the preference; a narrow band calls for a quantile, not a fixed value
Human review of the exemplars against their clusters, and for summarization a content-overlap measure against reference summaries. Cluster purity is meaningless without labels.
- A narrow similarity band, which makes the cluster count hypersensitive to the preference
- Quadratic memory on a large corpus; use a nearest-neighbour sparse similarity or sample first
- Treating the exemplar as the whole story of its cluster when the cluster is broad
Computer Vision
ViableRepresent each image or frame by an embedding, compute similarity as cosine or negative squared Euclidean distance between embeddings, and let the exemplars be real images. The method was originally demonstrated on faces and on gene-expression data for exactly this reason: a cluster centre that is itself an image can be shown to a person.
- Keyframe selection from video, where the chosen frames must be actual frames
- Choosing a representative subset of an image collection for labelling or for a gallery cover
- Grouping detected faces or objects and showing one real example per group
The exemplar is a real image, which makes the output reviewable by a person in a way an averaged embedding is not. The similarity should come from a learned embedding rather than raw pixels, where squared Euclidean distance mostly measures lighting and position. As elsewhere, the method is quadratic, so it suits a collection of thousands of items rather than millions, and the cost of choosing the preference badly is a gallery with either one picture or every picture.
- Use a learned embedding, not pixels, so the similarity reflects content
- Normalize embeddings if using cosine, and keep the same normalization when assigning new images to the existing exemplars
- Deduplicate near-identical frames first, since adjacent video frames dominate the similarity table
Coverage of the collection by the chosen exemplars — the mean similarity between each item and its exemplar — and human review of whether the exemplars are visually distinct from one another.
- Adjacent frames being near-identical, so the exemplars cluster in time rather than in content
- Pixel-space similarity, which clusters by brightness
- A preference chosen once and reused after the embedding model changes
What it costs to run
Dominated by memory: three n-by-n matrices resident at once, at roughly a gigabyte each in double precision at ten thousand points, and a few hundred iterations each costing O(n²). Float32 halves the footprint at a precision cost, and a sparse nearest-neighbour similarity removes the quadratic entirely for large inputs. Computing the similarity table itself is O(n²·d) for vectors and can exceed the message passing in cost when the similarity is something expensive such as time-warped distance.
Cheap and genuinely usable, which is a real contrast with spectral clustering: a new point is assigned to the exemplar it is most similar to, an O(exemplars) lookup with no refit. It is a nearest-exemplar rule, not a rerun of the algorithm — the new point cannot change which points are exemplars — so the assignment rule and the fitted clustering agree on the training data and can differ slightly beyond it.
Batch and periodic, and triggered by drift rather than the calendar: serve nearest-exemplar assignment in between, and refit when the average similarity to the assigned exemplar falls or the count of poorly-covered points rises. A refit can change the exemplar set, so keep cluster identity keyed to something stable rather than to the exemplar’s index.
- Track the mean similarity between each point and its exemplar — the coverage — and the share of points below a floor; falling coverage means the exemplars no longer represent the data
- Track the number of exemplars across refits; a jump means the preference has drifted relative to the similarity scale
- Track iterations to convergence and whether the run converged at all; a rising count is early warning of oscillation
- Compare exemplar sets across refits by overlap, and map clusters between runs by exemplar identity or membership overlap, never by cluster index
- Non-convergence can return nothing or a partial result without raising: always check the convergence flag and the exemplar count before using the labels
- The preference is in the units of the similarity, so changing the similarity function, the embedding model, or the normalization silently changes what the same preference means
- A distance passed where a similarity is expected selects the most outlying points as exemplars, and the output looks perfectly plausible
- Quadratic memory is hit without warning as the dataset grows between runs; the matrices, not the algorithm, are what fail
- The exemplar set can change when a few points are added, so any pipeline keyed on an exemplar’s identity breaks at refit
Pros and cons, in context
Needs only a similarity table, not coordinates, a metric, or symmetry
Sequences under edit distance, documents under cosine, series under correlation, directed costs that are not equal both ways: all are valid input. That opens problems where k-means is undefined because there is no mean, and where spectral clustering would need a symmetric affinity.
The number of clusters is not an input
The preference sets the price of a cluster and the count follows. It is a trade rather than an escape — the preference still has to be chosen — but it is chosen on the similarity scale, can be derived from the data, and can differ per point, which a fixed k cannot express.
Cluster centres are real, inspectable data points
The exemplar of a group of documents is a document, of a group of series is a series, of a group of customers is a customer. That makes the result reviewable by someone who cannot read a centroid vector, and it is the property that makes the method usable for summarization and representative selection.
No dependence on initialization
Every point starts as an equal candidate, so there is no random start and no restarts to run. The same input and settings give the same answer, which is a concrete operational advantage over k-means.
New points can be assigned without a refit
A new point goes to the exemplar it is most similar to. It is an approximation to rerunning the algorithm, but it exists, which spectral clustering does not offer.
Quadratic memory and quadratic time per iteration
Three dense n-by-n matrices cap the method near ten thousand points. A sparse neighbour similarity pushes it further, and costs whatever the pruning discards. This decides most affinity-propagation-versus-k-means arguments before cluster quality is discussed.
It can oscillate and has no convergence guarantee
Damping and iteration limits turn this into a tuning problem, and the failure is easy to misread as bad data. A result produced after hitting the iteration cap is not a result.
The preference is a hidden k with a worse unit
It controls the cluster count, but on the scale of the similarity and non-linearly, so reaching a target of ten clusters means scanning. A single value also cannot serve clusters of very different density, which is a real limit on realistic data.
Exemplars exclude elongated or ring-shaped clusters
Each cluster is represented by one member and each point joins the exemplar it is most similar to, so the clusters are compact around a leader. Chains and manifolds are cut into pieces — spectral clustering and DBSCAN are what to reach for there.
No optimality guarantee for a problem that has exact solvers
The underlying objective is facility location, where exact and provably approximate methods exist. On small and medium instances a solver beats this on both quality and certainty, and the method should be judged against one.
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""Affinity propagation - the message updates, transcribed literally.
S[i][k] says how well point k would serve as the exemplar for point i. The
diagonal S[k][k] is the PREFERENCE: k's own prior suitability to be an
exemplar, and the only thing that decides how many clusters emerge.
Two messages travel between every ordered pair of points:
R[i][k] responsibility, i -> k : how much i wants k as its exemplar,
measured against the best alternative i has.
A[i][k] availability, k -> i : how appropriate it would be for i to
pick k, given how much support k has from everyone else.
Taken literally, the max over rivals makes each responsibility O(n) and a
whole sweep O(n^3). That is the price of reading the formula line by line.
"""
import random
NOISE_SCALE = 1e-12 # relative jitter that breaks exact ties between similarities
def affinity_propagation(S, damping=0.5, max_iterations=500, stable_for=15, seed=0):
n = len(S)
rng = random.Random(seed)
# Exact ties leave the messages no way to break symmetry; a tiny relative
# jitter is enough, and does not change which exemplars win.
noisy = [
[S[i][k] + NOISE_SCALE * (abs(S[i][k]) + 1.0) * rng.random() for k in range(n)]
for i in range(n)
]
R = [[0.0] * n for _ in range(n)]
A = [[0.0] * n for _ in range(n)]
previous = None
unchanged = 0
for iteration in range(max_iterations):
# r(i,k) <- s(i,k) - max over k' != k of [ a(i,k') + s(i,k') ]
for i in range(n):
for k in range(n):
best_rival = float("-inf")
for rival in range(n):
if rival != k:
best_rival = max(best_rival, A[i][rival] + noisy[i][rival])
fresh = noisy[i][k] - best_rival
R[i][k] = damping * R[i][k] + (1.0 - damping) * fresh
# a(i,k) <- min(0, r(k,k) + sum over i' not in {i,k} of max(0, r(i',k)))
# a(k,k) <- sum over i' != k of max(0, r(i',k))
for i in range(n):
for k in range(n):
if i == k:
fresh = sum(max(0.0, R[other][k]) for other in range(n) if other != k)
else:
support = R[k][k] + sum(
max(0.0, R[other][k])
for other in range(n)
if other != i and other != k
)
fresh = min(0.0, support)
A[i][k] = damping * A[i][k] + (1.0 - damping) * fresh
# k is an exemplar when a(k,k) + r(k,k) is positive. The decision, not
# the message values, is what is watched for convergence.
exemplars = tuple(k for k in range(n) if A[k][k] + R[k][k] > 0.0)
unchanged = unchanged + 1 if exemplars == previous else 0
previous = exemplars
if unchanged >= stable_for and exemplars:
break
if not previous:
raise ValueError("no exemplars emerged: preference too low, or no convergence")
labels = []
for i in range(n):
if i in previous:
labels.append(i)
else:
labels.append(max(previous, key=lambda k: S[i][k]))
return labels, list(previous)
def net_similarity(S, labels):
"""The objective: sum of s(i, c_i). Each exemplar's own diagonal entry is
its preference, so every extra cluster pays its price here."""
return sum(S[i][labels[i]] for i in range(len(S)))O(n^3) time per iteration from the literal max over rivals and the nested availability sums, over three lists of lists.
How this language actually wants it written.
"""Affinity propagation - typed, vectorized, validated, and loud on failure."""
from dataclasses import dataclass
import numpy as np
from numpy.typing import NDArray
Matrix = NDArray[np.float64]
Indices = NDArray[np.int64]
NOISE_SCALE = 1e-12
MIN_DAMPING = 0.5
class NotConvergedError(RuntimeError):
"""The exemplar set never settled within the iteration budget."""
@dataclass(frozen=True)
class Clustering:
exemplars: Indices # which points were chosen to lead
labels: Indices # (n,) the exemplar each point follows
n_iter: int
net_similarity: float # the objective, preferences included
def negative_squared_euclidean(X: Matrix) -> Matrix:
"""The default similarity for vectors. For anything else, build the table
yourself - nothing below assumes a metric, or even symmetry.
The scale of this table is the scale of the preference: change one and the
other's values mean something different.
"""
squared_norms = np.einsum("ij,ij->i", X, X)
distances = squared_norms[:, None] - 2.0 * (X @ X.T) + squared_norms[None, :]
return -np.maximum(distances, 0.0) # cancellation guard
def _responsibilities(S: Matrix, A: Matrix) -> Matrix:
"""r(i,k) = s(i,k) - best rival, where the rival is the row's largest
a+s - or its SECOND largest when k itself is the largest.
Keeping the top two per row turns the O(n) max over rivals into a lookup,
which is the whole difference between O(n^2) and O(n^3) per sweep.
"""
rows = np.arange(S.shape[0])
scores = A + S
best = scores.argmax(axis=1)
best_value = scores[rows, best]
scores[rows, best] = -np.inf
second_value = scores.max(axis=1)
fresh = S - best_value[:, None]
fresh[rows, best] = S[rows, best] - second_value
return fresh
def _availabilities(R: Matrix) -> Matrix:
"""a(i,k) = min(0, r(k,k) + sum of positive r(i',k) over i' not in {i,k})."""
positive = np.maximum(R, 0.0)
np.fill_diagonal(positive, np.diag(R)) # self-responsibility is not clipped
column = positive.sum(axis=0)
fresh = np.minimum(0.0, column[None, :] - positive)
np.fill_diagonal(fresh, column - np.diag(R)) # a(k,k) = sum of the positives
return fresh
def fit(
similarity: Matrix,
preference: float | None = None,
damping: float = 0.7,
max_iter: int = 500,
stable_for: int = 15,
seed: int = 0,
) -> Clustering:
"""Cluster from a similarity table. Raises ValueError on a malformed
table and NotConvergedError when the exemplars never settle."""
if similarity.ndim != 2 or similarity.shape[0] != similarity.shape[1]:
raise ValueError(f"similarity must be square, got shape {similarity.shape}")
if np.isnan(similarity).any():
raise ValueError("similarity contains NaN; use -inf to forbid a pairing")
if not MIN_DAMPING <= damping < 1.0:
raise ValueError(f"damping must lie in [{MIN_DAMPING}, 1), got {damping}")
n = similarity.shape[0]
S = similarity.astype(np.float64, copy=True) # never mutate the caller's table
if preference is None:
preference = float(np.median(S[~np.eye(n, dtype=bool)])) if n > 1 else 0.0
np.fill_diagonal(S, preference)
rng = np.random.default_rng(seed)
jittered = S + NOISE_SCALE * (np.abs(S) + 1.0) * rng.random(S.shape)
R = np.zeros_like(S)
A = np.zeros_like(S)
previous: NDArray[np.bool_] | None = None
unchanged = 0
for iteration in range(1, max_iter + 1):
R = damping * R + (1.0 - damping) * _responsibilities(jittered, A)
A = damping * A + (1.0 - damping) * _availabilities(R)
is_exemplar = (np.diag(A) + np.diag(R)) > 0.0
same = previous is not None and np.array_equal(is_exemplar, previous)
unchanged = unchanged + 1 if same else 0
previous = is_exemplar
if unchanged >= stable_for and is_exemplar.any():
break
else:
raise NotConvergedError(
f"no stable exemplar set after {max_iter} iterations - "
"raise damping, or check the preference against the similarity scale"
)
exemplars = np.flatnonzero(is_exemplar)
labels = exemplars[np.argmax(S[:, exemplars], axis=1)]
labels[exemplars] = exemplars # an exemplar leads itself
net = float(S[np.arange(n), labels].sum())
return Clustering(exemplars=exemplars, labels=labels, n_iter=iteration, net_similarity=net)What changed: The O(n³) transcription becomes O(n²) per sweep: the max over rivals is replaced by the largest and second-largest value of each row, and the availability sums by one column reduction with the diagonal handled as the special case it is. The loops over points disappear into array operations. The parts that were comments become code: the similarity is validated and never mutated, damping is checked against its legal range, the preference defaults to the median of the off-diagonal similarities, and — the important change — non-convergence now raises a named exception instead of returning whatever the last iteration produced, because a result at the iteration cap is not a result. Convergence watches the exemplar set rather than the message values, the quantity that actually matters. The result is a frozen dataclass carrying the net similarity, so the objective is reported rather than implied.
O(n^2) time per iteration and about six n-by-n float64 temporaries live per sweep, which is what limits it before the algorithm does.
Faster, and the justification for each change.
"""Affinity propagation - in-place sweeps over preallocated float32 buffers."""
import numpy as np
from numpy.typing import NDArray
Matrix = NDArray[np.float32]
NOISE_SCALE = 1e-5 # float32 has ~1e-7 resolution; jitter below that vanishes
def fit_fast(
similarity: Matrix,
preference: float,
damping: float = 0.7,
max_iter: int = 500,
stable_for: int = 15,
seed: int = 0,
) -> tuple[NDArray[np.int64], NDArray[np.int64], int]:
"""Returns (exemplars, labels, iterations). One contiguous float32 dtype
throughout, and every sweep writes into buffers allocated exactly once."""
S = np.array(similarity, dtype=np.float32, order="C", copy=True)
n = S.shape[0]
np.fill_diagonal(S, preference)
R = np.zeros_like(S)
A = np.zeros_like(S)
scratch = np.empty_like(S) # the only other n-by-n buffer
column = np.empty(n, dtype=np.float32)
rows = np.arange(n)
# Strided VIEWS of the diagonals - reading them allocates nothing.
s_diag = S.reshape(-1)[:: n + 1]
r_diag = R.reshape(-1)[:: n + 1]
a_diag = A.reshape(-1)[:: n + 1]
rng = np.random.default_rng(seed)
rng.random(size=S.shape, dtype=np.float32, out=scratch)
scratch *= NOISE_SCALE
scratch *= np.abs(S) + 1.0 # relative jitter, applied once
S += scratch
keep, take = np.float32(damping), np.float32(1.0 - damping)
previous = np.zeros(n, dtype=bool)
unchanged = 0
for iteration in range(1, max_iter + 1):
# --- responsibilities: one buffer holds A + S, then the new R ---
np.add(A, S, out=scratch)
best = scratch.argmax(axis=1)
best_value = scratch[rows, best]
scratch[rows, best] = -np.inf
second_value = scratch.max(axis=1)
np.subtract(S, best_value[:, None], out=scratch)
scratch[rows, best] = S[rows, best] - second_value
R *= keep
scratch *= take
R += scratch
# --- availabilities: the clipped copy of R becomes the new A ---
np.maximum(R, 0.0, out=scratch)
scratch.reshape(-1)[:: n + 1] = r_diag # diagonal is not clipped
scratch.sum(axis=0, out=column)
np.subtract(column, scratch, out=scratch)
np.minimum(scratch, 0.0, out=scratch)
scratch.reshape(-1)[:: n + 1] = column - r_diag
A *= keep
scratch *= take
A += scratch
is_exemplar = (a_diag + r_diag) > 0.0
unchanged = unchanged + 1 if np.array_equal(is_exemplar, previous) else 0
np.copyto(previous, is_exemplar)
if unchanged >= stable_for and is_exemplar.any():
break
else:
raise RuntimeError(f"no stable exemplar set after {max_iter} iterations")
exemplars = np.flatnonzero(is_exemplar)
labels = exemplars[np.argmax(S[:, exemplars], axis=1)]
labels[exemplars] = exemplars
return exemplars, labels, iterationWhat changed: The right-stage version allocates a handful of fresh n-by-n arrays every sweep; at the sizes where this method is used, that allocation and the memory traffic behind it are the cost. Here everything is one contiguous float32 dtype — halving memory and bandwidth — and each sweep writes into three buffers allocated once: R, A, and a single scratch array that holds A + S, then the new responsibilities, then the clipped responsibilities, then the new availabilities, in turn. The row maximum and runner-up come from that one buffer instead of building a masked copy, the damping update is expressed as in-place scalings, and the diagonals are strided views rather than copies. What did not change is the algorithm and its ceiling: it is still three dense matrices, and past a few tens of thousands of points the answer is a sparse neighbour graph, not a tighter loop.
Pre-allocate output arrays and use in-place operations
R, A and one scratch matrix are allocated once and every sweep writes into them through out= arguments, so no n-by-n array is created inside the loop.
Trade-off: The single scratch buffer is reused for four different meanings in one sweep, so reordering any two lines silently reads stale data — correctness now lives in line order rather than in variable names.
Ensure contiguous memory layout and a single dtype
Everything is C-contiguous float32, which halves the memory and the bandwidth against float64, and makes the strided diagonal views and the row reductions single sequential passes.
Trade-off: Column sums over up to n terms in float32 lose precision at large n, and the tie-breaking jitter must be scaled to float32 resolution or it disappears — which brings back the oscillation it exists to prevent.
Fuse adjacent operations so intermediates are never materialized
The row maximum, its runner-up and the damped update all come out of the one A + S buffer, instead of a masked copy, a subtracted copy and a damped copy being built separately.
Trade-off: The fused sequence no longer reads like the two update formulas, which makes it hard to check against the derivation and hard to modify — and it has no sparse path, so the O(n^2) memory ceiling stays.
O(n^2) time per iteration with four n-by-n float32 buffers resident in total and zero per-iteration allocation. Illustrative, not a measured benchmark.