Caller-driven e-means training — em

Lifecycle: experimental. These functions may change in a later release.

They run e-means one batch at a time, for data that will not fit in memory or arrives in pieces, and mirror upstream faiss's begin/accumulate/update/end_epoch/finalize, with emeans() built on the same internals.

em_begin() creates a training state from the total row count and the starting centers. Its $recipe gives the batch size (batch_size) and the rows of each epoch (epoch_rows).
For each epoch e, the caller shuffles its rows and delivers the first recipe$epoch_rows[e] of them in consecutive batches of recipe$batch_size rows; the last batch of an epoch keeps the remainder. Each batch goes through em_accumulate() and then em_update(), and may reach em_accumulate() in several pieces that are all assigned against the same centers, with only the statistics of the call that completes the batch accepted by em_update().
em_end_epoch() closes each epoch.
em_finalize() returns the fit.

Batch rules, as upstream enforces them: every batch has exactly recipe$batch_size rows except the last of each pass; each pass delivers exactly recipe$epoch_rows[i] rows (the first epoch_rows[i] of a fresh shuffle); and calls come in the order above, em_accumulate() before em_update(), em_end_epoch() after the pass's last batch.

Batch sizes, epoch row counts and call order are checked. A rejected call changes nothing, so the loop can continue; an error after a call has started changing the state leaves it unusable, and later calls ask for a new state.

Usage

em_begin(
  n,
  centers,
  nepochs = 5,
  batch_law = c("auto", "throughput", "quality"),
  batch_size = NULL,
  alpha0 = NULL,
  spherical = FALSE,
  unused_threshold = NULL,
  chunk_size = NULL,
  centroid_chunk_size = NULL,
  device = "cpu",
  threads = NULL,
  verbose = FALSE
)

em_accumulate(state, batch)

em_update(state, stats)

em_end_epoch(state)

em_finalize(state)

em_centers(state)

# S3 method for class 'em_state'
print(x, ...)

Arguments

n: Total number of rows in the data, which may exceed R's integer range.
centers: A k x d numeric matrix of starting centers. In spherical mode they are normalised to unit length.
nepochs: Number of passes over the data; may be fractional.
batch_law: Scaling law that picks the batch size when batch_size is NULL. "auto" uses "quality" when d <= 16 and "throughput" otherwise.
batch_size: Rows per batch, or NULL to use batch_law.
alpha0: Initial EMA step in (0, 1], or NULL for the step law.
spherical: Cluster unit vectors by inner product, leaving normalisation of the data to you.
unused_threshold: Revival threshold: clusters whose mass is below this fraction of the mean are revived. NULL means 0.1, or 0 (no revival) when d <= 2.
chunk_size, centroid_chunk_size: Rows and centers searched at once, or NULL for the defaults. They trade memory for speed and, with the "direct" kernel, leave the result untouched.
device: Where to compute: "cpu" (the default) runs the C++ kernels in double precision; "mps" (Apple GPUs), "cuda" or "cuda:<index>" (NVIDIA GPUs, experimental), "torch:cpu" or a torch_device run the optional torch backend in float32, as described under "Devices" in emeans().
threads: Number of threads, or NULL for getOption("emeans.threads", 2L).
verbose: Print progress messages. The header is printed at the first call after em_begin().
state: A state from em_begin().
batch: A numeric matrix or data frame holding the batch's rows (or part of them), with d columns. On a torch device it may also be a 2-D torch_tensor, which is cast to float32 and not checked for NA.
stats: The result of the em_accumulate() call that completed the batch.
x: An em_state.
...: Ignored.

Value

em_begin(): an em_state, an environment whose $recipe is the training recipe.
em_accumulate(): an em_stats list with n (rows of the current batch so far), complete (whether the batch is full) and labels (1-based nearest-center labels of the rows just passed, in the last epoch only; otherwise NULL).
em_update(), em_end_epoch(): the state, invisibly.
em_finalize(): an "emeans" object with cluster = NULL and no sums of squares, since the caller keeps the last-epoch labels and the data; calling it again returns the same object.
em_centers(): the current centers as a k x d matrix (unit directions in spherical mode), or the final centers after em_finalize().

Examples

set.seed(1)
x <- matrix(rnorm(400), ncol = 2)
st <- em_begin(nrow(x), x[1:3, ], nepochs = 2, threads = 1)
B <- st$recipe$batch_size
for (e in seq_along(st$recipe$epoch_rows)) {
  ord <- sample.int(nrow(x))[seq_len(st$recipe$epoch_rows[e])]
  for (s in seq(1, length(ord), by = B)) {
    rows <- ord[s:min(s + B - 1, length(ord))]
    em_update(st, em_accumulate(st, x[rows, , drop = FALSE]))
  }
  em_end_epoch(st)
}
fit <- em_finalize(st)
fit$centers
#>             [,1]       [,2]
#> [1,] -0.01504195 -0.9480183
#> [2,]  1.03210137  0.4709296
#> [3,] -0.70823776  0.7415547