Skip to content
LAKER
pip install laker
v0.5.0 · PyTorch

Attention kernels, solved at scale.

LAKER is a PyTorch library for regularised attention-kernel regression with a learned preconditioner. It reduces kernel-system condition numbers by up to three orders of magnitude and converges in a near size-independent number of PCG iterations — turning expensive spectrum-cartography fits into something you can run in seconds.

Get started Read the paper View source
1000×
condition-number reduction
vs Jacobi preconditioner
100k+
samples, single node
GPU-accelerated PCG
10.38dB
masked RMSE on UCF-50K
vs 20.73 baseline
UCF-50K · MAP 01231 10.38 dB −100 −40 RSSI · dBm
Reconstructed RF field · 256×256 grid LAKER · Nyström-100
Citation
Tao & Tan, 2026 · arXiv:2604.25138
arXiv 2604.25138 License MIT Python 3.10 → 3.13 Build CI green Reproducibility paper · scalable

The problem

Attention is expressive.
Solving it is expensive.

Regularised attention-kernel regression problems look like a routine linear system — and that is exactly the trap. The kernel G = exp(E ET) explodes in condition number as the data grows, and vanilla PCG grinds to a halt long before your model converges.

LAKER turns that linear system into a tractable one. A learned, data-dependent preconditioner built by a shrinkage-regularised Convex-Concave Procedure (CCCP) slashes κ(G + λI) by up to three orders of magnitude — and the convergence of PCG becomes near size-independent.

See the architecture arXiv:2604.25138
Figure 1

Condition number vs sample size

10⁹ 10⁶ 10³ 10⁰ 10⁻³ CONDITION NUMBER κ(G + λI) 10² 10³ 10⁴ 10⁵ 10⁶ SAMPLE SIZE n ~10⁹ ~10³ ~1000× κ reduction Jacobi preconditioner LAKER · CCCP preconditioner
10⁶ → 10³
κ reduction on the UCF-50K kernel system at n = 50,000.
~constant
PCG iterations to 1e-6 tolerance stay bounded as n grows.
2 min
Reproduces the paper's Section V experiment end-to-end on CPU.

Capabilities

A complete toolkit
for attention-kernel
regression.

Every primitive you need — kernel approximation, learned embeddings, streaming updates, hyperparameter search — composes into a single coherent estimator. No glue code. No compromises.

01 flagship

Learned CCCP preconditioner

A factored, shrinkage-regularised Convex–Concave Procedure builds a data-dependent preconditioner in O(N_r³) — independent of problem size.

Embeddings E ∈ ℝⁿˣᵈ Encoder Attention G = exp(EEᵀ) + Low-rank CCCP shrinkage P ≈ (G + λI)⁻¹ factored PCG α* = (G+λI)⁻¹y ≈ constant iters
02 solver

Size-independent PCG

Preconditioned Conjugate Gradient converges in a near-constant number of iterations as n grows. No more O(n³) blow-ups.

10⁻⁶ tol LAKER · 14 iters Jacobi · >300 iters ‖r‖/‖r₀‖ PCG ITERATION
03

Seven kernel operators

Exact, Nyström, Fourier, sparse k-NN, SKI grid, spectrum shaping, and a two-scale hybrid. Matvec drops from O(n²) to O(n·r).

exact nystrom fourier neighbors grid · SKI spectrum hybrid · two-scale Nyström × Fourier blend
04

End-to-end learning

Learn the encoder end-to-end via backprop through the kernel operator. Bilevel hyperparameter learning via implicit differentiation.

Encoder E = φ(x) Kernel G = exp(EEᵀ) Solve PCG → α* Predict f(x) = Gα ∇_φ L · adjoint through K
05

Predictive variance

Exact posterior variance via batched PCG; closed-form for Fourier features via the Woodbury identity. Calibrate with NLL + a calibration penalty.

μ(x) ±2σ(x)
06

sklearn-compatible

fit / predict / score · get_params / set_params · __sklearn_clone__. Drop into scikit-learn pipelines and meta-estimators with no glue code.

fit · predict · score get_params · set_params __sklearn_clone__ Pipelines GridSearchCV 
# Compose with scikit-learn out of the box
from sklearn.pipeline import Pipeline
from laker import Laker

pipe = Pipeline([
    ("scaler", StandardScaler()),
    ("laker", Laker(embed_dim=10, lam=1e-2)),
])

pipe.fit(X_train, y_train)
y_pred = pipe.predict(X_test)

Architecture

One pipeline.
Five composable stages.

LAKER is a single Python class with five named steps. Each step is independent, inspectable, and replaceable. You can swap the kernel, swap the encoder, change the preconditioner — and still ship the same sklearn-compatible Laker estimator.

01 · ENCODE φ : ℝᵈ → ℝᵏ 02 · KERNEL G = exp(E Eᵀ) 03 · CCCP P P ≈ (G + λI)⁻¹ 04 · SOLVE PCG → α* 05 · PREDICT μ(x), σ²(x) INPUT · x ∈ ℝⁿˣᵈ y ∈ ℝⁿ · measurements μ̂(x*) · σ̂²(x*) α* ∈ ℝⁿ · coefficients learned · data-dependent O(n · r) matvec · r ≪ n ≈constant iters
01 Encode

Learned embeddings

A small encoder maps raw inputs x ∈ ℝⁿˣᵈ into an embedding E ∈ ℝⁿˣᵏ. Drop in a custom torch module (CNN, transformer) — the rest of the pipeline does not care.

02 Kernel

Exponential attention

The kernel matrix G = exp(E Eᵀ) is the heart of attention. We approximate it with seven operators — exact, Nyström, RFF, sparse k-NN, SKI, spectrum, hybrid — at O(n·r) cost.

03 Precondition

CCCP factored P

A shrinkage-regularised Convex-Concave Procedure factorises P ≈ (G + λI)⁻¹. Cost is O(N_r³) per construction step and amortised across PCG iterations.

04 Solve

PCG, fast

Preconditioned Conjugate Gradient converges in a near-constant number of iterations, regardless of n. Optionally batch across multiple right-hand sides.

05 Predict

Posterior + variance

Recover α* = (G + λI)⁻¹ y, then form predictions and exact posterior variances — closed-form for RFF, batched PCG otherwise.

Objective
minα ‖G α − y‖² + λ α⊤ G α
Solve with PCG. Speed it up with a learned preconditioner. Backprop through the kernel to learn the encoder. See the paper.

Quick start

From pip install
to a fitted model
in three lines.

LAKER ships with a single public entry point — Laker — and the full sklearn API. Install, import, fit, predict. No bespoke training loop. No boilerplate.

pip install laker Browse the source
Python 3.10 → 3.13
Single source of truth, modern toolchain.
CPU, CUDA, MPS — same API
PyTorch-native device dispatch.
310 tests · MIT licensed
Real behavioural assertions, no mocks.
fit.py 
# Install: pip install laker
import torch
from laker import Laker

# A small synthetic radio field on a 100×100 grid.
n = 1_000
x = torch.rand(n, 2) * 100.0
y = torch.sin(x[:, 0] / 50) + torch.cos(x[:, 1] / 50)

# One line. CCCP preconditioner is automatic.
model = Laker(embed_dim=10, lam=1e-2, device="cuda")
model.fit(x, y)

# Predict anywhere on the field.
x_test = torch.rand(2000, 2) * 100.0
y_pred = model.predict(x_test)
y_var  = model.variance(x_test)        # exact posterior variance

print(f"R² = {model.score(x, y):.3f}")        # → R² = 0.998
print(f"κ = {model.condition():.1e}")           # → κ ≈ 1.2e+03
$ pip install laker $ python fit.py R² = 0.998 · κ ≈ 1.2e+03
CLI · laker fit shell 
# Train from the command line
laker fit --locations x_train.pt --measurements y_train.pt --output model.pt

# Predict on a held-out set
laker predict --model model.pt --locations x_test.pt --output y_pred.pt

Benchmarks

Numbers that earn their place.

LAKER is evaluated on the paper's synthetic scene and on the full UCF-50K corpus — 50,000 ray-traced radio maps with the complete masked 256×256 grid. Reproducible end-to-end via examples.scalable.

Headline RMSE
10.38dB
Masked RMSE on UCF-50K — nyström-100, λ = 1e-2.
vs 20.73 dB baseline
Maps validated
50,000
Full corpus, complete masked 256×256 grid, resumable run.
20260801T114917Z
κ reduction
1000×
CCCP preconditioner vs Jacobi on the LAKER kernel system.
10⁶ → 10³
Paper experiment
~2min
examples.paper on a CPU; reproduces Section V.
arXiv:2604.25138
Figure 2

Masked RMSE on UCF-50K

Lower is better · validation on the complete masked 256×256 grid

25 dB 20 dB 15 dB 10 dB 5 dB 20.73 baseline 20.73 cross-map 14.10 nystrom-200 10.38 nystrom-100 λ = 1e-2 winner 9.92 exact (ref) −10.4 dB vs baseline
The CCCP-preconditioned Nyström-100 fit lands at 10.38 dB, ~50% below the per-scene mean baseline of 20.73 dB, and within ~5% of the exact dense reference solve.
Figure 3

PCG iterations vs n

Tolerance 1e-6 · n up to 50k

800 600 400 200 0 PCG ITERATIONS 1k 5k 10k 25k 50k SAMPLE SIZE n Jacobi · > 800 LAKER · ~12
Iterations to 1e-6 tolerance stay near-constant with the CCCP preconditioner — they grow roughly linearly without it.
All numbers produced by the pinned snapshots under benchmarks/ and examples/. Run python -m benchmarks.reproducible to regenerate.

Choosing a kernel

Pick the right one.
We pick it automatically.

LAKER ships with seven kernel operators. The right one depends on the size of n, the embedding dimension, and the structure of the data. These thresholds come from the UCF-50K sweep under outputs/scalable/.

exact

Exact

n ≤ 5,000

Memory ≈ 100 MB at float32; the exact path is fastest and most accurate.

recommended
nystrom

Nyström

5k < n ≤ 50k

Low-rank matvec; matches `exact` to within 5% relative error on UCF-50K.

fourier

Fourier

n > 50,000

Cheaper than Nyström at very large n; some accuracy loss on fast-growing exponential kernels.

grid

Grid · SKI

embed_dim ≤ 4

Product grid is only practical in low dimensions — use Nyström or RFF beyond that.

Switch kernels with one argument.

kernel_type is the only knob. Everything else — the encoder, the preconditioner, the solver — stays the same. 

# Try every kernel in a few lines
for k in ["exact", "nystrom", "fourier", "grid"]:
    m = Laker(embed_dim=10, kernel_type=k, lam=1e-2)
    m.fit(x_train, y_train)
    print(f"{k:>8s} · R² = {m.score(x_test, y_test):.3f}")

Get started

Fit your first attention
kernel regression in five minutes.

Install from PyPI, read the paper, browse the source. The library is MIT-licensed and the benchmarks are reproducible end-to-end.

pip install laker Read the paper Examples
your_first_fit.py 3 lines 
from laker import Laker
import torch

model = Laker(embed_dim=10)
model.fit(x_train, y_train)

print(model.predict(x_test))
$ python your_first_fit.py tensor([0.314, -0.812, …])