🤖 machine learning

LAPLEX: The FFT of Learnable Laplace Kernels

LAPLEX introduces a class of exact, trainable Laplace-kernel operators that achieve FFT-like scaling to enable data-adaptive global interactions and high-dimensional covariance modeling without the storage costs of dense matrices, effectively separating expressivity from memory requirements.

Original authors: Łukasz Struski, Hanna Blazhko, Piotr Kubaty, Jacek Tabor

Published 2026-05-26

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Łukasz Struski, Hanna Blazhko, Piotr Kubaty, Jacek Tabor

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to organize a massive library with billions of books. In the world of modern artificial intelligence, "books" are just numbers (data points), and "organizing" means performing complex math to find patterns.

Usually, when AI tries to do this on a huge scale, it faces a frustrating choice:

The Rigid Librarian (Fixed Geometry): Like a standard library catalog, it's fast and exact, but the shelves are fixed. You can't move a book to a better spot based on what you're looking for. (This is like the FFT or Fourier Transform).
The Flexible but Heavy Librarian (Dense Parameters): This librarian can rearrange the shelves however they want to fit the books perfectly, but to do so, they need a massive warehouse to store the instructions for every single book. If the library gets too big, the warehouse runs out of space, and the system crashes. (This is a standard Dense Matrix).
The Approximate Librarian (Low-Rank/Sketching): This librarian tries to save space by only remembering the "general vibe" of the books or using a random guessing game. It's fast and fits in the warehouse, but it loses the exact details and can't learn the best arrangement.

LAPLEX is a new kind of librarian that breaks this trade-off. It claims to be the "Goldilocks" solution: it is exact, it is trainable (it can learn the best arrangement), and it is lightweight (it doesn't need a massive warehouse).

Here is how it works, using simple analogies:

1. The Secret: Learning the "Coordinates" Instead of the "Map"

Imagine you have a giant map of a city with millions of streets.

The Old Way: To know the distance between every pair of streets, you would need to write down a number for every single pair. For a city with 1 million streets, that's a trillion numbers. You can't store that.
The LAPLEX Way: Instead of writing down every distance, LAPLEX just writes down a list of "anchor points" (like major landmarks). It then uses a clever mathematical rule (the Laplace Kernel) to calculate the distance between any two points based on how close they are to these landmarks.

Because the rule is so simple, LAPLEX doesn't need to store the trillion numbers. It only needs to store the list of landmarks. But here's the magic: it can learn where to put those landmarks. If the data suggests that "Street A" and "Street B" are often related, LAPLEX learns to move its landmarks so that A and B end up close together in its internal logic.

2. The "Scan" Trick: Doing Math Without the Heavy Lifting

The paper claims that even though LAPLEX acts like a giant, dense map, it can do the math incredibly fast.

Think of it like a conveyor belt in a factory.

The Dense Way: To calculate the result for one item, you have to stop, look at every other item in the warehouse, and do a calculation. This gets slower and slower as the warehouse grows.
The LAPLEX Way: It sorts the items on the conveyor belt first. Then, it uses a "scan" technique. Imagine a worker walking down the line who only needs to remember the running total of what they've seen so far. Because the math of the Laplace kernel is exponential (things get smaller very fast as you move away), the worker can calculate the final result for the whole line by just looking at the "prefix" (what came before) and the "suffix" (what comes after).

This allows LAPLEX to handle dimensions up to 1 billion (10⁹) on modern computers, a size where the "Dense" method would immediately run out of memory.

3. What It Actually Does (The Experiments)

The paper tests this "librarian" in three specific scenarios to prove it works:

The "Flattened Image" Test: They took high-resolution photos (millions of pixels), flattened them into a single long line of numbers (ignoring the 2D grid), and tried to model the data.
- Result: A standard "low-rank" model (the approximate librarian) failed to see that pixels next to each other in the original photo should be related; it just saw noise. LAPLEX, however, learned to put its "landmarks" in a way that preserved the shapes of the images (like a beach or a palm tree) even though it never saw the image as a grid. It learned the geometry of the data without being told what a grid is.
The "Classification Head" Test: They tried to replace the final layer of a neural network (the part that decides "is this a cat or a dog?") with LAPLEX.
- Result: Using only a tiny fraction of the parameters (like 1-5% of the usual size), LAPLEX performed almost as well as the massive, full-size network. It learned to route information much better than standard "low-rank" shortcuts.
The "Speed" Test: They compared how fast it runs on a supercomputer (GPU).
- Result: LAPLEX is roughly 65 times faster than the standard method for forward passes and 430 times faster for training (forward + backward) on large datasets. It also uses 100 times less memory.

4. Why This Matters (In Simple Terms)

The paper argues that for a long time, we thought we had to choose between speed/size and accuracy/flexibility.

If you wanted speed, you had to use fixed rules (like FFT) or random guesses.
If you wanted to learn the best rules, you had to use so much memory that you couldn't run it on big data.

LAPLEX says: "You don't have to choose." By treating the coordinates of the data as learnable variables rather than fixed slots, it creates a system that is:

Exact: It doesn't guess; it calculates the true math.
Trainable: It adapts to the specific data it sees.
Efficient: It fits on hardware that would otherwise crash.

Summary Analogy

Imagine you are trying to organize a massive party where millions of guests need to find their seats.

Dense Matrix: You print a giant seating chart with every guest's name next to every other guest's name. It's perfect, but the paper pile is too heavy to carry.
Low-Rank/Sketch: You just tell guests to sit in "Zone A" or "Zone B" randomly. It's light, but the seating is messy and inefficient.
LAPLEX: You give the guests a few "magnetic anchors" (like VIP tables). You teach the guests to sit based on how close they are to these anchors. The anchors move around during the party to find the best spots. You don't need a giant chart; you just need the list of anchors. The guests find their seats instantly, the arrangement is perfect, and you can carry the whole plan in your pocket.

The paper concludes that this approach allows AI to handle "dense" interactions (where everything talks to everything) at scales that were previously impossible, without sacrificing the ability to learn from the data.

Technical Summary: LAPLEX – The FFT of Learnable Laplace Kernels

1. Problem Statement

Modern machine learning systems increasingly operate in regimes where ambient dimensions reach millions or billions. In these high-dimensional settings, the bottleneck is often the wall-clock cost and memory footprint of linear operations, such as applying dense matrices, computing Gram operators, and mixing features.

Existing approaches to scaling linear algebra typically sacrifice one of three desirable properties: exactness, learnability, or near-linear scaling.

Low-rank methods (e.g., LoRA) retain dense, basis-dependent interactions but suffer from an $O(nk)$ memory footprint because they must store coefficients for the input basis, even if the output rank is compressed.
Randomized sketches and random features (e.g., CountSketch, Random Fourier Features) offer scalability but act on fixed indexing or approximate the kernel, trading exactness for efficiency. They cannot adapt how coordinates are matched.
Sparse methods can break hardware regularity and are difficult to train end-to-end.
Classical fast transforms (e.g., FFT) are exact and cheap ( $O(n \log n)$ ) but are fixed to a specific grid geometry and cannot learn the coordinate system.

The core limitation identified is basis-dependence: standard operators parameterize the matrix values directly, requiring storage proportional to the input dimension. The paper argues for a shift to index-level operators that can learn the input geometry itself without storing a dense matrix.

2. Methodology: LAPLEX

The authors introduce LAPLEX, a class of exact, trainable operators based on the Laplace kernel. A LAPLEX layer is implicitly defined by learnable coordinate anchors rather than explicit matrix weights.

2.1 The Operator

The core primitive is the Laplace-kernel matrix defined by two sets of learnable anchors, $a \in \mathbb{R}^n$ and $b \in \mathbb{R}^k$ :
$\text{LAPLEX}(a, b)_{ij} = \exp(-|a_i - b_j|)$
In a natural limiting regime (temperature $t \to 0$ ), this operator encodes all permutations of the coordinate set. By adding phase parameters $\phi, \psi$ , the operator can be extended to:
$\text{LAPLEX}(a, \phi; b, \psi)_{jk} = \exp(-|a_j - b_k|) \cdot \cos(\phi_j - \psi_k)$
This allows the operator to function as a dense, trainable matrix while using only $O(n)$ parameters (the anchors).

2.2 Exact Scan Algebra

The key algorithmic contribution is an exact, fully parallel $O(n \log n)$ algorithm for matrix-vector and Gram products that never materializes the dense matrix.

Symmetric Case ( $a=b$ ): After sorting the anchors, the kernel matrix decomposes into lower and upper triangular parts. Each part can be computed via a prefix-scan and suffix-scan recurrence. For example, the lower triangular product $y_i = \sum_{j \le i} \exp(a_j - a_i)x_j$ satisfies a recurrence $y_{i+1} = x_{i+1} + \exp(a_i - a_{i+1})y_i$ , which is evaluated in $O(n)$ work and $O(\log n)$ parallel depth.
General Case ( $a \neq b$ ): The algorithm aggregates contributions from $b$ -anchors into intervals defined by sorted $a$ -anchors. Local contributions are summarized into vectors $u$ and $v$ , which are then propagated globally using the symmetric scan primitive.
Gram Operators: The same ordering structure enables exact computation of weighted Gram matrices $M = A \text{diag}(D) A^\top$ in $O(n^2 + k \log n)$ time, avoiding the naive $O(n^2 k)$ cost. This is crucial for second-order optimization and kernel methods.

2.3 Expressivity and Limits

The paper establishes that LAPLEX is a continuous family interpolating between different operator classes:

Zero-Temperature Limit: As the temperature parameter approaches zero, the smooth kernel converges to a discrete equality-routing matrix, recovering CountSketch and arbitrary permutations.
Universality: With diagonal scalings and sparse replication, the family is universal for arbitrary rectangular dense matrices.
Relation to FFT: If anchors are frozen to a uniform grid, LAPLEX becomes a Toeplitz matrix, and the algorithm reduces to Toeplitz-FFT. If the kernel is replaced by a finite spectral expansion, it becomes Random Fourier Features (RFF). LAPLEX retains exactness and differentiability with respect to the anchors, unlike RFF which suffers from unbounded gradient variance due to the heavy-tailed Cauchy spectral density of the Laplace kernel.

3. Key Contributions

Identification of Basis-Dependence: The paper motivates index-level, permutation-aware operators as a principled alternative to low-rank methods, which cannot escape $O(nk)$ memory footprints because they do not compress the input basis.
Learnable Geometry: It demonstrates that the Laplace-kernel matrix encodes all coordinate permutations in a natural limiting regime while using only $O(n)$ parameters.
Exact, Scalable Algorithms: The authors provide an exact, fully parallel $O(n \log n)$ algorithm for matrix-vector and Gram products using $O(n)$ memory. This enables dimensions up to $10^9$ on modern GPUs without materializing the matrix.
Empirical Validation: The paper shows that LAPLEX matches FFT-class scaling while functioning efficiently as a drop-in structured layer. It outperforms dense baselines in memory and speed, and surpasses low-rank and random feature baselines in accuracy and gradient stability.

4. Experimental Results

4.1 Scaling and Performance

GPU Scaling: On NVIDIA H100 and RTX 2080, LAPLEX scales smoothly to $n = 2^{20}$ (approx. 1 million elements), whereas dense baselines run out of memory (OOM) at $n = 2^{14}$ .
Speed: At the largest size where both methods run ( $n=2^{14}$ ), LAPLEX is $\sim 65\times$ faster on the forward pass and $\sim 430\times$ faster on forward+backward compared to dense matmul.
Memory: LAPLEX uses $\sim 100\times$ less peak memory than the dense baseline.
Numerical Accuracy: Despite running in float32, LAPLEX exhibits lower relative error against a float64 reference than both dense float32 matmul and Toeplitz-FFT. This is attributed to the scan algorithm's ability to keep exponential scale factors separate from mantissas, avoiding the $O(n \epsilon)$ error accumulation of dense summation.

4.2 High-Dimensional Density Modeling

The authors trained a Gaussian density model on flattened DIV2K images (dimension $3 \cdot 10^6$ ) using an exact Gram operator.

Setup: Compared a LAPLEX factor (rank 1000, $\sim 14$ M parameters) against a dense low-rank baseline (rank 4, $\sim 15$ M parameters).
Result: The dense baseline collapsed to uniform noise due to insufficient rank. The LAPLEX model learned coherent spatial color structures (contiguous regions, silhouettes) despite operating on flattened vectors with no convolutional bias. The test log-likelihood was $4.5\times$ higher than the dense baseline.

4.3 Structured Projections and Classification

ImageNet10 Projections: In a frozen projection setting (flattened $150,528$-dim images to $m \in \{25, \dots, 1000\}$ ), LAPLEX was more stable than CountSketch and Random Projections, and became the best performer at larger $m$ .
Trainable Heads: When used as a classification head for Swin-T and MaxViT-T, LAPLEX matched or exceeded the accuracy of unconstrained dense projections at the same output dimension, using $\Theta(m+d)$ parameters instead of $\Theta(md)$ . At low ranks, LAPLEX significantly outperformed low-rank factorization (LoRA-style).

5. Significance and Claims

The paper claims that LAPLEX separates expressivity from storage cost. It behaves like a dense trainable matrix but is represented and applied through a small structured set of parameters (anchors).

Separation of Expressivity and Storage: Dense geometry can be learned without storing a dense matrix. This enables data-adaptive global interactions in regimes where ordinary dense layers are infeasible.
Exactness vs. Approximation: Unlike Random Fourier Features or Nyström methods, LAPLEX provides exact kernel application and Gram computation. It avoids the value/gradient trade-offs of random features, particularly the unbounded variance in anchor gradients caused by the Cauchy spectral density.
Learnable Fast Transforms: The work demonstrates that fast transforms do not have to live on fixed coordinate systems. By parameterizing a dense operator through learnable Laplace anchors, the geometry itself becomes a differentiable parameter.

6. Limitations

The authors acknowledge the following limitations:

Dimensionality: The current construction uses one-dimensional Laplace anchors and sorting-based algorithms. Extending this exactness to higher-dimensional anchor geometries and richer kernels is an open mathematical direction.
Systems Maturity: The implementation demonstrates the algorithmic regime, but large-scale attention blocks, optimizers, and generative models will require specialized kernels and careful architecture design.
Scope: The method is currently limited to the Laplace kernel and its specific scan-based decomposition.

In summary, LAPLEX offers a principled alternative to low-rank and randomized methods, enabling exact, learnable, and scalable linear algebra in high-dimensional regimes by shifting from parameterizing matrices to parameterizing index interactions.