Skip to content

Support reusing pre-computed UMAP embeddings across runs (grid search, caching) #2504

Description

@pidefrem

Feature request

BERTopic always computes UMAP internally via _reduce_dimensionality(). While BaseDimensionalityReduction lets users skip UMAP entirely (#909, #2385), there's no way to reuse specific UMAP embeddings from a previous run.

Add a umap_embeddings parameter to fit(), fit_transform(), and transform() that accepts pre-computed UMAP embeddings, skipping the UMAP step when provided:

# Run UMAP once
topic_model = BERTopic()
topics, probs = topic_model.fit_transform(docs, embeddings=embeddings)
umap_embs = topic_model.umap_model.embedding_

# Reuse UMAP embeddings with different HDBSCAN settings
for min_size in [10, 15, 20, 30]:
    model = BERTopic(hdbscan_model=HDBSCAN(min_cluster_size=min_size))
    topics, probs = model.fit_transform(docs, embeddings=embeddings, umap_embeddings=umap_embs)

Motivation

Two common workflows are blocked:

  1. Hyperparameter grid search — same UMAP reduction, different HDBSCAN min_cluster_size. Currently UMAP (the most expensive step) re-runs for every HDBSCAN configuration.
  2. Caching across sessions — save UMAP embeddings to disk, reload them without re-running UMAP. Useful for large datasets where UMAP takes minutes.

BaseDimensionalityReduction doesn't solve these cases because it replaces the UMAP model entirely — the user loses access to the fitted UMAP for later use with .transform().

Your contribution

I can submit a PR that adds a umap_embeddings parameter to fit(), fit_transform(), and transform(). When provided, skip UMAP and use the given embeddings directly. The UMAP model is still stored (for .transform() on new data) but not called during that fit.

All new parameters default to None — existing code unchanged. Complementary to BaseDimensionalityReduction, not a replacement.

I've already been prototyping this in my fork, so I can open a PR quickly if this direction works for you.


Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions