[KongchangAI]
· 2 min read· 1,277 words

Gaussian-Splat EBM: AUROC on Par with kNN Density Estimation, 3.9x Faster Inference

Gaussian-Splat EBM: AUROC on Par with kNN Density Estimation, 3.9x Faster Inference

Gaussian-splat EBM ties kNN density estimation on AUROC but scores faster, uses constant memory, and supports generation and composition.

A developer's Gaussian-splat EBM achieves AUROC nearly on par with classic kNN-density (0.9585 vs 0.9584) across 20 text-embedding anomaly detection tasks, but scores 3.9x faster with constant memory that doesn't grow with the corpus. The model represents the corpus as Gaussian splats on a sphere, aggregates nearest-neighbor energy via logsumexp, and supports both Rectified Flow and Langevin sampling, enabling generation, corpus composition, and incremental growth. A key failure—initial alpha calibration yielding AUROC 0.06—was fixed to 0.96 by setting alpha=5.0. CUDA graph capture speeds up 100-step sampling 12.9x and fused kernels 3.2-3.4x. Code is Apache-2.0 open source with reproducible benchmarks and a fixed semantic bug.

In the field of anomaly detection on text embeddings, classic density-based methods (such as kNN density estimation) have long held a dominant position. But a Gaussian-splat energy model (Gaussian-splat EBM) built by a developer achieves AUROC nearly on par with kNN-density across 20 tasks (0.9585 vs. 0.9584), while showing clear advantages in scoring speed, memory usage, and extensibility. The results come from a complete benchmark shared on Reddit, with code and reproducible experiments open-sourced.

Screenshot of the original Reddit post

Model Design: A Gaussian-Splat Energy Field on the Sphere

The model represents the entire corpus as a set of Gaussian splats in embedding space. Each splat is defined by three parameters: position (mu), sharpness (alpha), and weight (kappa). The energy function is defined as the negative of a logsumexp operation over the nearest k splats. This design allows the model to store the entire distribution with compact parameters, without retaining all original embeddings.

Sampling supports two methods: Rectified Flow and Langevin dynamics. According to the developer, in Phase-1 tests, Rectified Flow sampling is 29-39x faster than 200-step Langevin. This generative capability is something traditional methods like kNN-density do not offer.

Energy-based models (EBMs) are a class of generative models that assign a scalar value to each data point via an energy function, where lower energy indicates a sample is more likely to come from the true distribution. Unlike traditional probabilistic models, EBMs require no explicit normalization constant, and training and sampling often rely on Markov chain methods such as Langevin dynamics. The Gaussian-splat EBM in this article treats each splat as a local energy well in embedding space, aggregating the energy of nearest-neighbor splats via logsumexp, thereby representing the entire corpus distribution with compact parameters. This design enables the model to serve as a one-class classifier for anomaly scoring while retaining generative and compositional capabilities.

logsumexp is a smooth approximation of the maximum operation, defined as (\log \sum_i \exp(x_i)). In the energy function, taking the negative logsumexp over the energies of the nearest k splats is equivalent to taking a soft minimum of those splat energies, making the score sensitive to local density changes and numerically stable. Compared with a hard minimum, logsumexp is differentiable and avoids vanishing gradients, making it suitable for gradient-based sampling and optimization.

Rectified Flow is a generative sampling method based on ordinary differential equations (ODEs) that rectifies straight-line paths between the noise distribution and the data distribution to achieve high-quality sampling in few steps. Unlike Langevin dynamics, which requires hundreds of iterations, Rectified Flow typically needs only tens or even a few steps to generate samples, thus achieving the 29-39x speedup in this work. Langevin dynamics, by contrast, is a sampling method based on stochastic gradient descent that explores the distribution by injecting noise into the energy gradient; it converges to the target distribution in theory but is slow. Combining the two allows the model to both generate quickly and sample accurately.

Benchmark: A Statistical Tie, but Significant Operational Advantages

Using the same encoder, data splits, and random seeds, and fitting only on positive samples (no negative labels), the developer compared multiple one-class methods across 20 tasks. Key results are below:

MethodAUROC (20 tasks)Median time per document (ms)Memory usage
ebm-splats0.9585 ± 0.0180.45386 KB (constant)
knn-density0.9584 ± 0.0201.75802 KB (grows with corpus)
mahalanobis0.9523 ± 0.0205.26578 KB
gmm (diag)0.9439 ± 0.0230.38192 KB
centroid0.9324 ± 0.034<0.012 KB
iforest0.8828 ± 0.05118.71200 KB

In terms of quality, the difference between ebm-splats and kNN-density is only 0.0001, a statistical tie. But the operational differences are more important: CPU scoring is 3.9x faster, and memory complexity is O(kd) rather than O(nd), meaning memory does not grow with corpus size. In addition, the splat model supports generation (via Rectified Flow), composition of multiple corpora (adding energy fields), and incremental growth—three things classic methods cannot do simultaneously.

To be clear, when labels are available, supervised classifiers beat all one-class methods. This comparison only concerns method ranking in the unlabeled setting.

A Key Lesson from a Failure

The most valuable finding came from a failure. The developer initially used alpha = ln(cluster_size) as a "natural" calibration, which resulted in an AUROC of only 0.06—systematically inverted. The reason: for small clusters, alpha approaches 0, the distance term disappears, and scoring degenerates into ranking by splat weight. Positive samples near small, specific splats ended up with higher energy than negative samples.

After fixing alpha at 5.0 (making alpha/tau much greater than 1), AUROC jumped from 0.06 to 0.96. The developer warns: if you fit splats without contrastive training, you must be wary of this pitfall.

AUROC (Area Under the Receiver Operating Characteristic Curve) is a commonly used evaluation metric in anomaly detection, measuring a model's ability to distinguish positive from negative samples across thresholds. An AUROC of 0.5 indicates random guessing, while 1.0 indicates perfect separation. When AUROC is below 0.5, it means the model systematically inverts the scores of positive and negative samples—that is, positive samples receive higher anomaly scores than negative samples. In this article, the failed alpha calibration led to an AUROC of just 0.06, a classic case of score inversion, indicating that model parameter settings have a decisive influence on ranking direction.

GPU Kernel Optimization: CUDA Graph Capture and Fused Kernels

Since scoring speed is critical, the developer also wrote GPU kernels (Triton + CUDA graphs) for the hot path and ran equally honest benchmarks.

Full-loop CUDA graph capture speeds up 100-step Langevin sampling by 12.9x (from 660 ms to 51 ms), and graph replay is bit-for-bit exact with eager execution. The fused energy/scoring kernel brings a 3.2-3.4x speedup.

Not every attempt succeeded: one kernel ended up 15-40x slower and has been recorded as a rejected approach; the KNN kernel only wins on cold lookups, and the system automatically routes around it.

CUDA graph capture is a technique that packages a series of GPU kernel launches and their parameters into a single graph structure, replaying the entire computation sequence with a single submission, thereby eliminating CPU overhead and driver latency from repeated kernel launches. For iterative algorithms containing many small kernels (such as Langevin sampling), graph capture can significantly improve throughput. In this article, full-loop CUDA graph capture reduces 100-step sampling from 660 ms to 51 ms, and graph replay is bit-for-bit exact with eager execution, demonstrating the reliability of this technique in achieving substantial speedups without changing numerical results.

Open Source Code and Reproducibility

The repository is open-sourced under the Apache-2.0 license, installable via pip install ebm-splats, with CI passing and reproducible benchmarks. The developer also documented and fixed a semantic bug called find_neighbors: the function had selected the farthest splats since its inception, and the two code paths were silently inconsistent.

Repository: https://github.com/schwabauerbriantomas-gif/EBM-splats

Benchmark methodology and the full table are in benches/ONECLASS_DUEL_RESULTS.md; kernel comparisons are in benches/bench_headtohead.py.

Why Keep This "Merely Tied" Model

Although it ties with kNN-density on AUROC, the developer chose to keep the model because of operational advantages: constant memory, faster scoring, and generative and compositional capabilities. These properties make the model more scalable and flexible in real systems. For scenarios requiring incremental updates, multi-corpus composition, or generative sampling, the Gaussian-splat EBM offers capabilities that classic methods cannot cover.

This work also demonstrates the value of honest benchmarking: by publicly sharing failures and rejected approaches, it provides reusable lessons for the community.

#Gaussian-splat EBM#anomaly detection#kNN density estimation#AUROC#CUDA graph capture#Langevin sampling#Rectified Flow#text embeddings
Share:

Related articles