Fit e-means clustering

emeans() clusters the rows of x into k groups, much as stats::kmeans() does, but learns from small random batches instead of the whole data set at once. Each batch pulls its nearest centers part of the way towards it and the pull weakens over a few passes, until a closing pass puts every center at the mean of its points. The authors of e-means report a lower loss than Lloyd's k-means in about a tenth of the passes (Bessa et al. 2026). A final sweep then labels every row by its nearest center, so the fit carries the same cluster, size and sums of squares as a stats::kmeans() fit.

Most fits need only x and k. The remaining arguments either adjust the training recipe, which emeans() otherwise chooses from the size of the data, or decide where and how fast the work runs.

Usage

emeans(
  x,
  k,
  nepochs = 5,
  batch_law = c("auto", "throughput", "quality"),
  batch_size = NULL,
  alpha0 = NULL,
  centers = NULL,
  spherical = FALSE,
  unused_threshold = NULL,
  max_points_per_centroid = NULL,
  keep_cluster = TRUE,
  chunk_size = NULL,
  centroid_chunk_size = NULL,
  device = "cpu",
  threads = NULL,
  verbose = FALSE
)

Arguments

x: The data, one row per observation: a numeric matrix or an all-numeric data frame without NA, NaN or infinite values. On a torch device it may also be a 2-D torch_tensor (see "Devices").
k: Number of clusters, at most the number of rows used for training.
nepochs: Number of passes over the data, 5 by default. It may be fractional, such as 2.5, but not so small that a pass would hold no rows.
batch_law: How the default batch size is chosen. "quality" picks smaller batches for a lower loss and "throughput" larger ones for speed, while "auto", the default, uses "quality" for data with at most 16 columns and "throughput" beyond that.
batch_size: Rows per batch. The default, NULL, lets batch_law choose.
alpha0: Starting step size, between 0 and 1, which sets how far a center moves towards a batch early on. The default, NULL, derives it from k, the batch size, the number of rows and nepochs.
centers: Optional k x d matrix of starting centers. By default emeans() picks k distinct rows of x at random.
spherical: If TRUE, cluster directions rather than points: rows are compared by inner product and the centers kept at unit length. Normalise the rows of x yourself first.
unused_threshold: Clusters holding less than this fraction of the average cluster's weight are revived by splitting a large cluster in two. The default, NULL, means 0.1, or no revival for data with one or two columns.
max_points_per_centroid: If set, train on at most max_points_per_centroid * k rows chosen at random; the rows left out get NA in cluster. The default, NULL, trains on every row.
keep_cluster: If FALSE, skip the final sweep that labels every row, which saves roughly the time of one pass and the memory for the labels. The fit then has no cluster or sums of squares, and size counts the rows of the last training pass.
chunk_size, centroid_chunk_size: How many rows and centers are compared at once. They trade memory for speed and rarely need changing; NULL uses the defaults.
device: Where to compute: "cpu" (the default), "mps" for Apple GPUs, "cuda" or "cuda:<index>" for NVIDIA GPUs (experimental), "torch:cpu" or a torch_device. Every device other than "cpu" needs the torch package (see "Devices").
threads: Number of CPU threads. The default, NULL, uses getOption("emeans.threads", 2L).
verbose: If TRUE, print the recipe and a progress line for each pass.

Value

An object of class "emeans", which is a list with:

cluster: the nearest final center for each row of x, as an integer vector matching predict(fit, x). Rows left out by max_points_per_centroid are NA.
centers: the k x d matrix of final centers.
totss: the total sum of squares about the overall mean.
withinss: for each cluster, the sum of squared distances from its rows to its center.
tot.withinss: sum(withinss).
betweenss: totss - tot.withinss.
size: the number of rows in each cluster.
iter: the number of batch updates.
history: a data frame with one row per pass, holding epoch (passes so far), rows (rows seen in the pass), obj (the sum of squared distances from each row to its center when it was seen), time (seconds since the start), imbalance (1 when the clusters are equal in size, larger otherwise) and nsplit (clusters revived).
recipe: the settings used, including nepochs, batch_size, alpha0, epoch_rows and device.
n_splits: the total number of clusters revived.
spherical: whether the fit is spherical.
call: the matched call.

The sums of squares use squared Euclidean distance and cover the rows used for training. With keep_cluster = FALSE the fit has no cluster, totss, withinss, tot.withinss or betweenss, and size counts the rows of the last training pass. In spherical mode the centers have unit length, so the sums of squares compare with totss only when the rows of x do too.

Devices

On "cpu", emeans runs C++ code in double precision. Every other device runs the same algorithm in float32 through torch and returns an ordinary "emeans" object, and predict() then assigns new data on the same device. set.seed() still controls every random draw, but the results differ slightly from the CPU because float32 rounds differently, and on CUDA they can even differ from one run to the next.

The CUDA path has not been tested on NVIDIA hardware, so please report what you find at https://github.com/coatless-rpkg/emeans/issues. A torch_tensor passed as x is not checked for NA, NaN or infinite values, since that would mean copying it back from the device.

Speed

threads sets the number of CPU threads, up to any limit in the OMP_THREAD_LIMIT environment variable. R for macOS is built without OpenMP, so on a Mac emeans uses one thread.

On the CPU, emeans uses a BLAS-based kernel when R is linked to an optimised BLAS and a simpler direct kernel otherwise. Set options(emeans.kernel = "direct") when you need results that are identical to the last bit whatever the chunk sizes.

Differences from upstream

emeans ports the e-means code in faiss (contrib/e_means), with these deliberate differences:

R's random number generator and set.seed() replace numpy's RandomState and the seed argument, so a run does not match Python for the same seed.
The CPU path uses double precision throughout, where upstream uses float32 for the model and float64 for the objective.
Cluster labels start at 1, and rows left out of training get NA rather than -1.
A tie for the nearest center goes to the lowest index.
The split-collision step moves to the next representable double, where upstream uses the next float32.
emeans() ends with the extra sweep that labels every row by its nearest final center and adds the sums of squares. Upstream returns the labels of the last training pass, which em_accumulate() still gives.

References

Bessa S, Mazaré P-E, Szilvasy G, Touvron H, Picard D, Douze M, Jégou H (2026). e-means: outperforming k-means (in one-tenth the epochs). Preprint. https://arxiv.org/abs/2611.09501

Examples

set.seed(1)
x <- rbind(matrix(rnorm(200, sd = 0.3), ncol = 2),
           matrix(rnorm(200, mean = 3, sd = 0.3), ncol = 2))
fit <- emeans(x, 2, threads = 1)
fit$centers
#>            [,1]        [,2]
#> [1,] 0.03266621 -0.01134242
#> [2,] 3.00890206  3.01548056

# Assign new points to the nearest center
predict(fit, rbind(c(0, 0), c(3, 3)))
#> [1] 1 2