Skip to content

BIRCH clustering - #347

Open
paulsullivanjr wants to merge 1 commit into
elixir-nx:mainfrom
paulsullivanjr:birch-clustering
Open

BIRCH clustering#347
paulsullivanjr wants to merge 1 commit into
elixir-nx:mainfrom
paulsullivanjr:birch-clustering

Conversation

@paulsullivanjr

@paulsullivanjr paulsullivanjr commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Adds Scholar.Cluster.BIRCH.

BIRCH makes a single pass over the data, compressing it into a
height-balanced tree of Clustering Features — additive (n, linear_sum, squared_sum) triples that summarize a group of points without storing
them. The leaf centroids are then reduced to :num_clusters, and each
input point is labelled by its nearest subcluster. 2000 points compress
to ~34 subclusters before the global step runs.

The CF-tree is built host-side in Elixir since insertion is sequential;
the labelling and global steps are tensor work.

model = Scholar.Cluster.BIRCH.fit(x, num_clusters: 3, key: key)
model.labels
Scholar.Cluster.BIRCH.predict(model, new_x)

Options mirror scikit-learn's defaults: :threshold (0.5),
:branching_factor (50), :num_clusters (3), plus :key,
:num_runs, :max_iterations for the global step.

Validation. The CF-tree is independent of the global step, so it was
differential-tested against scikit-learn 1.9.0 across 650 random cases
(varying n, dimensionality, branching factor, threshold; half on an
integer grid to force ties): 248/250 and 398/400 exact agreement
on leaf centroids at 1e-9. Residuals are last-ULP boundary decisions in
deep trees. 18 tests plus doctests, including three sklearn reference
tests and a regression test pinning subcluster ordering on splits.

One deviation: the global step uses Scholar.Cluster.KMeans rather
than agglomerative clustering, following SpectralClustering's
precedent. Scholar.Cluster.Hierarchical can't be used — it reorders
dendrogram rows without remapping the clade ids stored inside them, so
labels_list crashes on 25/25 random 50-point datasets. Pre-existing
and unrelated to BIRCH; left alone here, happy to file separately.

Not included: n_clusters=None, transform/2, partial_fit.

@josevalim

josevalim commented Aug 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Thank you @paulsullivanjr! Unfortunately this pull request does a lot of work outside of Nx, which may:

  1. Indicate this algorithm does not fit Scholar
  2. We need to refactor the algorithm

@RicardoSantos-99 can you take a look if this could mostly written as defn, as in your previous PRs?

Scholar.Cluster.Hierarchical can't be used — it reorders dendrogram rows without remapping the clade ids stored inside them, so labels_list crashes on 25/25 random 50-point datasets. Pre-existing and unrelated to BIRCH; left alone here, happy to file separately.

I am submitting a fix for that, thanks for filing!

@RicardoSantos-99

RicardoSantos-99 commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Nice work on the sklearn parity, @paulsullivanjr.

@RicardoSantos-99 can you take a look if this could mostly written as defn, as in your previous PRs?

@josevalim, Not with the CF-tree in place, I think. defn wants static shapes, and the tree is dynamic in node count, subcluster count and depth. Preallocating for the worst case means roughly one node per sample, which gives up the memory advantage that makes BIRCH worth using.

It also isn't buying much here. The tree exists to avoid scanning every subcluster, but scanning them all is one matvec for us.

Without it, the first phase is a single while over the samples with preallocated buffers, same shape as parallel_nearest_neighbor/3 in Hierarchical. The global step, labelling and predict are already fine.

The cost is the sklearn parity, and :branching_factor stops meaning anything.

@RicardoSantos-99

RicardoSantos-99 commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Not included: n_clusters=None, transform/2, partial_fit.

@paulsullivanjr You mention partial_fit as out of scope, which is fair for a first PR, though it's the piece that usually justifies BIRCH over KMeans, and PCA already has incremental_fit/2 as a precedent. Worth noting the flat version makes it easier rather than harder: the state is just the CF buffers and a counter, which is close to resumable already.

@RicardoSantos-99

RicardoSantos-99 commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

@paulsullivanjr here's what I had in mind, concretely, in case it's useful.

What comes out is everything from insert_root/4 down through closest_index/2, so the tree build and the CF helpers around it. global_labels/3, assign_labels/3, predict/2, the struct and the options schema stay as they are.

In its place, phase 1 becomes one while over the samples. The CF buffers are preallocated and a counter tracks how many rows are live:

while {ls, ss, counts, centers, m = 0, i = 0, x}, i < n do
  p = x[i]
  active = Nx.iota({max_m}) < m
  dist = Nx.select(active, Nx.sum((centers - p) ** 2, axes: [1]), inf)
  j = Nx.argmin(dist)

  n_j = counts[j] + 1
  ls_j = ls[j] + p
  ss_j = ss[j] + Nx.sum(p * p)
  radius2 = ss_j / n_j - Nx.sum((ls_j / n_j) ** 2)
  fits? = m > 0 and radius2 <= threshold ** 2

  idx = Nx.select(fits?, j, m)
  ls = Nx.put_slice(ls, [idx, 0], Nx.reshape(Nx.select(fits?, ls_j, p), {1, d}))
  # same for ss, counts and centers

  {..., m + Nx.select(fits?, 0, 1), i + 1, x}
end

Nx.put_slice/3 takes a dynamic start index, which is what lets the same line either merge into row j or open a new row at m without the shapes changing. You'd want a :max_subclusters option to size the buffers, defaulting to the sample count, and :branching_factor would go.

Worth being clear that this is still sequential over the samples, and nothing makes it otherwise, since each one depends on the ones before it. The gain is that the loop compiles into a single program instead of reading scalars back per sample. parallel_nearest_neighbor/3 in Hierarchical is in the same position, and dbscan_inner/2 has a comment explaining why they swapped a sequential DFS for label propagation.

The awkward part is the tests. The reference values won't survive, since a flat scan picks a different subcluster than the descent does. I'd rewrite them around properties instead: the radius bound holds for every subcluster, every point lands in its nearest one, results stay order dependent. Happy to help with that.

Two unrelated things I noticed while reading. The moduledoc says agglomerative clustering on line 9 and KMeans further down. And assign_labels/3 could use pairwise_squared_euclidean, since the square root doesn't change the argmin.

@josevalim

josevalim commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

@krstopro so I am thinking this one may not be a good fit for Scholar after all. :(

@krstopro

krstopro commented Aug 6, 2026

Copy link
Copy Markdown
Member

@krstopro so I am thinking this one may not be a good fit for Scholar after all. :(

@josevalim Quite possible, lemme have a look...

@krstopro

krstopro commented Aug 6, 2026

Copy link
Copy Markdown
Member

@krstopro so I am thinking this one may not be a good fit for Scholar after all. :(

@josevalim I agree, there are a lot of operations being performed outside of Nx which makes it not suitable for Scholar.

@paulsullivanjr I don't see a convenient way of doing this within defn. Do you maybe see one? If not, I would suggest closing this one.

@paulsullivanjr

paulsullivanjr commented Aug 6, 2026 via email •
edited
Loading

Copy link
Copy Markdown
Contributor Author

@paulsullivanjr

paulsullivanjr commented Aug 6, 2026 via email •
edited
Loading

Copy link
Copy Markdown
Contributor Author

@paulsullivanjr

paulsullivanjr commented Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

@paulsullivanjr here's what I had in mind, concretely, in case it's useful.

What comes out is everything from insert_root/4 down through closest_index/2, so the tree build and the CF helpers around it. global_labels/3, assign_labels/3, predict/2, the struct and the options schema stay as they are.

In its place, phase 1 becomes one while over the samples. The CF buffers are preallocated and a counter tracks how many rows are live:

while {ls, ss, counts, centers, m = 0, i = 0, x}, i < n do
  p = x[i]
  active = Nx.iota({max_m}) < m
  dist = Nx.select(active, Nx.sum((centers - p) ** 2, axes: [1]), inf)
  j = Nx.argmin(dist)

  n_j = counts[j] + 1
  ls_j = ls[j] + p
  ss_j = ss[j] + Nx.sum(p * p)
  radius2 = ss_j / n_j - Nx.sum((ls_j / n_j) ** 2)
  fits? = m > 0 and radius2 <= threshold ** 2

  idx = Nx.select(fits?, j, m)
  ls = Nx.put_slice(ls, [idx, 0], Nx.reshape(Nx.select(fits?, ls_j, p), {1, d}))
  # same for ss, counts and centers

  {..., m + Nx.select(fits?, 0, 1), i + 1, x}
end

Nx.put_slice/3 takes a dynamic start index, which is what lets the same line either merge into row j or open a new row at m without the shapes changing. You'd want a :max_subclusters option to size the buffers, defaulting to the sample count, and :branching_factor would go.

Worth being clear that this is still sequential over the samples, and nothing makes it otherwise, since each one depends on the ones before it. The gain is that the loop compiles into a single program instead of reading scalars back per sample. parallel_nearest_neighbor/3 in Hierarchical is in the same position, and dbscan_inner/2 has a comment explaining why they swapped a sequential DFS for label propagation.

The awkward part is the tests. The reference values won't survive, since a flat scan picks a different subcluster than the descent does. I'd rewrite them around properties instead: the radius bound holds for every subcluster, every point lands in its nearest one, results stay order dependent. Happy to help with that.

Two unrelated things I noticed while reading. The moduledoc says agglomerative clustering on line 9 and KMeans further down. And assign_labels/3 could use pairwise_squared_euclidean, since the square root doesn't change the argmin.

@RicardoSantos-99 Makes sense, better fit for Scholar than what I did, and agreed on moving the tests to properties.

Good catches on both, line 9 is a leftover from when the global step moved to KMeans and squared is right since the argmin is unchanged.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@paulsullivanjr @josevalim @RicardoSantos-99 @krstopro