Speed up step_word_embeddings()

[EmilHvitfeldt]

Note

This PR is a showcase of using {debrief} to drive a profiling-based optimization. Each iteration below came from profiling with profvis + debrief.

Origin

This started from the following prompt:

look the following code

library(recipes)
library(textrecipes)

embeddings <- tibble(
  tokens = c("the", "cat", "ran"),
  d1 = c(1, 0, 0),
  d2 = c(0, 1, 0),
  d3 = c(0, 0, 1)
)

sample_data <- tibble(
  text = c(
    "The.",
    "The cat.",
    "The cat ran."
  ),
  text_label = c("fragment", "fragment", "sentence")
)

rec <- recipe(text_label ~ ., data = sample_data) |>
  step_tokenize(text) |>
  step_word_embeddings(text, embeddings = embeddings)

rec |> prep() |> bake(sample_data)

I want to See if there are room for improvements. specifically in the step_word_embeddings part.

Use the debrief package like so:

library(profvis)
library(debrief)

# Profile some code
p <- profvis({
  # your code here
})

# Get help on available functions
pv_help()

# Start with a summary
pv_print_debrief(p)

you will need to profile this line

rec |> prep() |> bake(sample_data)

This is a iterative process. In a document write down each itteration that you take. Do a commit for each accepted change. you might need to generate new data for embeddings and sample_data to properly test the speed as this data set is quite small

Summary

Speeds up the bake() path of step_word_embeddings(), which was driven by tokenlist_embedding() in R/tokenlist.R. The original implementation sliced the embedding tibble by row and ran a dplyr group_by() + summarise_all() pipeline per document. This was rewritten over several profiling-driven iterations into matrix/sparse-matrix operations.

On a benchmark of 5,000 documents (10-30 tokens each), a 50,000-token vocabulary, and 100 embedding dimensions:

aggregation	baseline	this PR	speedup
sum	0.85 s	0.027 s	~31x
mean	0.85 s	0.027 s	~31x
min	0.85 s	0.325 s	~2.6x
max	0.85 s	0.301 s	~2.8x

Output is identical to before for all four aggregations, including documents with no matched tokens falling back to aggregation_default. All word_embeddings tests pass.

Iterations

The work was done as a sequence of profile-measure-change steps (full notes and benchmark script live in bench/, which is .Rbuildignored).

Baseline. Profiling showed nearly all the cost in tokenlist_embedding(): slicing a 100-column tibble by row (vec_slice, ~140 MB allocated) and the dplyr group/summarise machinery, plus heavy GC.

Iteration 1: matrix + rowsum. Convert the embeddings to a numeric matrix and slice rows from it (matrix slicing is far cheaper than tibble slicing), pass the aggregation name through instead of an opaque closure, and aggregate sum/mean with base rowsum(). → sum/mean ~13x; min/max ~2.2x (still using a per-group apply()).

Iteration 2: vectorise min/max. Replace the per-group apply() with a loop over the ~100 embedding columns, each aggregated by a single tapply() (far fewer iterations). → modest further gain on min/max.

Iteration 3: sparse incidence product for sum/mean. Re-profiling showed rowsum.default and the row-slice as the new hot spots. Replaced them with a single sparse-by-dense matrix product: a (n_doc x n_vocab) document-token incidence matrix times the dense embedding matrix yields per-document sums directly, in document order, with no row-slice and no grouped reduction. The embedding matrix stays dense here; only the incidence matrix is sparse, and it is sparse by construction (each document holds only a few of the vocabulary's tokens) regardless of the embedding values. Matrix is already a dependency. → sum/mean ~19-22x.

Iteration 4: precompute the dense matrix in prep(). as.matrix() was running on every bake(), allocating the whole ~40 MB embedding matrix each call and driving much of the GC. Moved the conversion into prep(): the prepped step now stores a dense emb_matrix plus an emb_tokens lookup vector instead of the embeddings tibble (no memory duplication, since an all-numeric matrix is no larger than the tibble), and tidy() reads the row count from emb_matrix. → sum/mean roughly halved again, ~31x overall.

Rejected experiments

For min/max I also tried a single flattened tapply over all columns at once (~2x slower) and an order() + duplicated(fromLast=) pick per column (~12% faster). The latter's gain was too small for the added complexity on a non-default path, so it was not kept.

Notes

• The default sum/mean path is now down to tokenization plus the intrinsic matrix multiply; there is little structural overhead left to remove.
• The full devtools::test() run segfaults inside an unrelated ngram C-code test under load_all; this reproduces on a clean main tree (pre-existing, not caused by this PR). The word_embeddings tests pass (40 PASS).

[EmilHvitfeldt]

The GHA failures are unrelated to this step