Really Fast Bayesian Linear Regression

I run contextual bandits in production as web services, and the model underneath is Bayesian linear regression, because it’s straightforward and powerful. A bandit on that model does two things. Every decision is a draw from the posterior and every reward is an update to it; once it’s deployed, those two calls are the whole compute bill. Optimizing them is worthwhile because their cost sets both what the model costs to run and how much model we can afford under a given latency budget.

I’d been making the code fast by minimizing algebraic operations, i.e. fewer factorizations, fewer solves, never forming a matrix you don’t need. Recently, I read about the roofline model¹ and wondered what it would say about my own library. The idea is that every routine is bound by one of two things, how fast the core can move bytes or how fast it can do arithmetic, and both are a one-minute measurement on the machine in front of you. Once you have them, every timing has a denominator, and you can tell a routine that’s done from one that’s merely faster than it was.

Everything below is one core of an i9-9940X, one thread, minimum of seven repeats, under OpenBLAS 0.3.31 and MKL 2026.1². Every version of the model was checked against a plain numpy implementation on 200,000 draws before anything was timed.

The floors

We rely on two measurements. The first is how fast one core can stream bytes, which depends on where the bytes live. Between the core and main memory sit three caches, each bigger and slower than the last: on this machine 32 KB of L1, 1 MB of L2 per core, and 19 MB of L3. A $p \times p$ precision matrix stored as one triangle is $4p^2$ bytes, so at $p = 256$ the triangle fits in L2, at $1{,}024$ in L3, and at $4{,}096$ it spills to main memory, and each level has its own speed:

Below that, at $p = 64$, the triangle is 16 KB and fits in L1, and the timings are dominated by the cost of calling numpy at all rather than by the memory system.

The second measurement is how fast the core does arithmetic when a library is really trying, which means a big dgemm (double, general matrix, matrix-matrix multiply): 89 GFLOP/s under OpenBLAS, 86 under MKL. dgemm does so much arithmetic on every byte it reads that the memory system never gets to be the bottleneck, so it’s the ceiling for any routine with the same property. Routines that do only a few operations per byte are bound by the first table instead.

With those two you can price the model before writing code. At $p = 1{,}024$:

The update floors assume the factor is updated in place rather than recomputed. Refactorizing is $np^2 + p^3/3$, which at 100 rows is 5.2 ms.

The math

Weights $\mathbf{w} \in \mathbb{R}^p$, noise precision $\beta$, prior $\mathcal{N}(\mathbf{w}_0, \boldsymbol{\Lambda}_0^{-1})$. In precision form³:

The posterior becomes the next prior, so a stream of batches is a stream of additions to $\boldsymbol{\Lambda}$ and one solve each. Each step has a routine with its name on it, in BLAS (the matrix-vector and matrix-matrix primitives) or LAPACK (the factorizations built on top of them):

• $\boldsymbol{\Lambda}_0 + \beta \mathbf{X}^\top \mathbf{X}$ is dsyrk (symmetric rank-k update, $\alpha \mathbf{A}^\top\mathbf{A} + \beta\mathbf{C}$), with the prior passed in as $\mathbf{C}$. It writes one triangle of the result, so it’s half the work of dgemm.
• $\boldsymbol{\eta}$ is dsymv (symmetric matrix-vector) for $\boldsymbol{\Lambda}_0 \mathbf{w}_0$, then dgemv (general matrix-vector) with beta=1 accumulating $\beta\mathbf{X}^\top\mathbf{y}$ onto it.
• The factorization is dpotrf (positive definite, triangular factorization, which is LAPACK for Cholesky), giving $\boldsymbol{\Lambda} = \mathbf{L}\mathbf{L}^\top$.
• The mean is two dtrsv calls: triangular solve, one vector.
• A draw is one more: if $\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ then $\mathbf{L}^{-\top}\mathbf{z}$ has covariance $\boldsymbol{\Lambda}^{-1}$, so a Thompson sample is $\mathbf{w}_N + \mathbf{L}^{-\top}\mathbf{z}$ and the covariance matrix never exists. Nothing calls multivariate_normal⁴.

 1def update(self, X, y):
 2    XF, xt = _fview(X)  # Fortran view, no copy
 3    eta = dsymv(1.0, self.Lam, self.w, lower=1)
 4    eta = dgemv(self.beta, XF, y, trans=1 - xt, beta=1.0, y=eta, overwrite_y=True)
 5    dsyrk(self.beta, XF, trans=1 - xt, lower=1, beta=1.0, c=self.Lam, overwrite_c=True)
 6    self.L, info = dpotrf(self.Lam, lower=1)
 7    self.w = dtrsv(self.L, dtrsv(self.L, eta, lower=1), lower=1, trans=1)
 8
 9
10def draw(self, Xq):
11    z = self.rng.standard_normal(len(self.w))
12    w = dtrsv(self.L, z, lower=1, trans=1, overwrite_x=True)  # L⁻ᵀ z
13    w += self.w
14    return _predict(Xq, w)  # dgemv against the query rows

The precision matrix lives as a single triangle for its whole life, since dsyrk, dsymv and dpotrf all read one triangle and ignore the other. The catch is that Lam @ v is silently wrong, because numpy reads the garbage on the other side of the diagonal. I have had that bug, from one product that didn’t go through dsymv.

The draw

The precision factor is a covariance square root

A draw is $\mathbf{w}_N$ plus a square root of the covariance times white noise. Taken literally, that means forming $\boldsymbol{\Lambda}^{-1}$ and taking its Cholesky: factor $\boldsymbol{\Lambda}$, solve for $\boldsymbol{\Lambda}^{-1}$, factor that – 13, 69 and 14 ms at $p = 1{,}024$ – for something whose floor is reading 4 MB once. But the Cholesky factor of $\boldsymbol{\Lambda}$ is a square root of the covariance too, just upside down: $\mathbf{L}^{-\top}\mathbf{z}$ has covariance $\boldsymbol{\Lambda}^{-1}$. That factor has already been prepared by the mean solve, so a draw is one triangular solve against it and nothing else.

dtrsv for one right-hand side

dtrsm (triangular solve, matrix of right-hand sides) solves against a whole matrix of columns at once; dtrsv solves against one. BLAS calls these level 3 and level 2, matrix-matrix and matrix-vector, and the level 3 routines are built around reusing the triangle across many columns. For one right-hand side you want dtrsv. My code called dtrsm for everything, because a Thompson draw was just size=1 on the same path as a thousand draws.

OpenBLAS’s dtrsm first copies the triangle into a cache-friendly layout, which pays off when many right-hand sides get to reuse the copy and is pure overhead for one.

The draw is at its floor

Across sizes, dtrsv runs at whatever speed the cache level holding the triangle can deliver:

The draw at $p = 1{,}024$ – normals, solve, projection, and the Python around it – is 190 µs, of which the solve is about 140. There’s no routine that reads a triangle less than once.

The update

Refactorizing after 100 rows at $p = 1{,}024$ takes 14.9 ms. The arithmetic alone, $np^2 + p^3/3$ at dgemm speed, would take 5.2.

f2py copies C-ordered arrays

scipy.linalg.blas is a generated wrapper (f2py) around Fortran routines, and Fortran stores a matrix column by column where numpy’s default is row by row. Hand a wrapper a numpy-default array, C-ordered in numpy’s terms, and it silently copies it into column order before every call. That copy is a transpose, done one element at a time, and it runs nowhere near a straight memcpy:

Whether it matters depends on how much work the routine does per byte. dsyrk does about $p$ flops per element of $\mathbf{X}$, so at 1,024 columns the copy is 6% of the call, and at 256 it’s more than twice the routine:

The middle row is how I found this. My update built $\mathbf{X}$ weighted by row, in C order like everything numpy produces, and passed it straight to dsyrk. At 5,000 rows the plain X.T @ X update is 8.2 ms, and mine was 32.7 – four times slower than the naive code it was supposed to be an optimization of, on the same routine with the same flops. With the copy gone it’s 7.4. The fix is no big deal: a C-contiguous array’s transpose is Fortran-contiguous, and every BLAS routine takes a transpose flag, so hand over X.T and flip the flag.

1def _fview(A):
2    """(F, trans) with F Fortran-contiguous and A == op(F)."""
3    if A.flags.f_contiguous:
4        return A, 0
5    if A.flags.c_contiguous:
6        return A.T, 1
7    return np.asfortranarray(A), 0

Don’t accidentally use two OpenBLAS thread pools

The numpy and scipy wheels each ship their own OpenBLAS, and each starts its own pool of worker threads sized to the logical CPU count – here, two pools of 28 threads on 14 cores, each under the impression that the machine is theirs. Workers busy-wait for a while after a job, holding a core while they check for more work, before going to sleep⁵. So a call into numpy’s pool followed by one into scipy’s has the second pool fighting the first, which is still holding cores for work that’s gone elsewhere. $\mathbf{X}^\top\mathbf{X}$ followed by its Cholesky at $p = 1{,}024$:

The mixed column is what you get by default, and at 28 threads it pays 3.5x for crossing libraries. It gets worse as the matrices get smaller, because the work gets shorter and the busy-waiting doesn’t: a thousand draws at $p = 64$ that cross from numpy into scipy take 1.4 ms on one thread and 24 ms on 28. And the threads weren’t doing much of use anyway – the only operation that got faster with more of them was the $p = 1{,}024$ factorization. I decided it’s best to keep the hot path entirely inside scipy.linalg, which has everything I need, and to pin OPENBLAS_NUM_THREADS=1 in the serving process.

dpotrf

dsyrk on 100 rows runs at 81 GFLOP/s, 1.1x its floor, so the rest of the update is dpotrf:

The gap has two parts, and comparing the libraries separates them. A Cholesky factorization works through the matrix a block of columns at a time. Factoring the block itself is matrix-vector work that can’t run at dgemm speed; only the update it makes to the rest of the matrix can. That slow slice is a fixed cost per block, so it’s a shrinking fraction of the whole as $p$ grows, which is why both libraries close in on the floor at 4,096. At 1,024, MKL’s 1.7x is roughly what the slow slice costs, and OpenBLAS’s 2.9x is the slow slice plus matrix-matrix code running at a third of its own dgemm speed. The update spends 12 of its 14.9 ms here, and none of it is reachable from Python.

Which triangle?

dsyrk and dsymv don’t care. dpotrf does: OpenBLAS is within noise between upper and lower at 1,024 and about 8% faster on lower at 256. MKL’s upper is twice as fast as its lower at 1,024 – 6.7 ms against 12.7. I use lower because I measured under OpenBLAS. Under MKL that choice costs 6 ms per update. I don’t really know why they differ, but I guess it’s an important reminder to benchmark on your target architecture and system library setup.

One row at a time

If dpotrf is the wall, the way past it is to stop calling it. Refactorizing is $p^3/3$ flops whether the batch has a thousand rows or one, and an online learner mostly sees one. The copy and thread-pool findings above apply to every path; this one just has no Cholesky in it.

Updating the factor

Keep $\mathbf{U}$ with $\boldsymbol{\Lambda} = \mathbf{U}^\top \mathbf{U}$, stack the new rows under it, and take a QR:

The same $\mathbf{Q}^\top$ applied to $[\mathbf{d};\ \sqrt{\beta}\,\mathbf{y}]$ with $\mathbf{d} = \mathbf{U}\mathbf{w}$ gives the new $\mathbf{d}$; a draw is $\mathbf{U}^{-1}(\mathbf{d} + \mathbf{z})$. This is the square-root information filter⁶. LAPACK has a routine for exactly this shape, a triangle on top of a rectangle: dtpqrt (triangular-pentagonal QR; the T is for the compressed form of $\mathbf{Q}$ it returns instead of the full matrix), at $2np^2$ flops, with dtpmqrt (multiply by that $\mathbf{Q}$) to apply it to the right-hand side. Conveniently, both are in scipy.linalg.lapack.

1def update(self, X, y):
2    B = np.multiply(X, self.sb, order="F")  # √β X
3    self.U, V, T, _ = dtpqrt(0, 16, self.U, B, overwrite_a=True, overwrite_b=True)
4    b = np.multiply(y, self.sb).reshape(-1, 1, order="F")
5    self.d, _, _ = dtpmqrt(
6        0, V, T, self.d, b, side="L", trans="T", overwrite_a=True, overwrite_b=True
7    )

dtpqrt is quite slow

Counting operations, a one-row update should beat refactorizing by 170x at $p = 1{,}024$. It beats it by 5.8, 2.5 ms against 14.2. I used to file that kind of shortfall under “constants” and move on. Against the floor it’s a different story: touching a 4 MB triangle once, read and write, is 224 µs, and dtpqrt takes ten times that.

My first explanation was that OpenBLAS ships the unmodified reference Fortran for this routine, which it does, and which is a satisfying explanation because it blames someone else. But MKL takes 2.1 ms. The problem is the algorithm. The triangle is stored column by column, so consecutive elements of a row are 8 KB apart. dtpqrt folds the new row in one column at a time, and each step touches a row of the triangle⁷. Memory arrives in 64-byte lines, so every element it touches drags in eight doubles and uses one. Blocking contains this by walking the row only across a block of 16 columns and doing the rest 16 rows at a time, which uses the whole line. That’s the difference between 65 ms unblocked and 2.5 ms at a block size of 16, and it’s also as far as blocking can go, because the inner walk is still in the wrong direction.

dtplqt

dtplqt is the same routine with LQ in place of QR, which is QR of the transpose. Transpose the whole construction: keep $\mathbf{L}$ with $\boldsymbol{\Lambda} = \mathbf{L}\mathbf{L}^\top$, put the new row beside it as a column, $[\mathbf{L},\ \sqrt{\beta}\,\mathbf{x}]$, and the inner routine now applies each reflection down a column of the triangle, which is the direction it’s stored in. scipy doesn’t wrap it, but the symbol is in the OpenBLAS that scipy bundles, and in MKL⁸.

At 256 the triangle is in L2, where jumping around in memory costs almost nothing extra, and the two routines take the same time. At 4,096 it’s in main memory, every wasted line is paid in full, and walking columns is 3.6 times faster. What’s left at every size is the blocking itself: the routine builds and applies a $16 \times 16$ compressed $\mathbf{Q}$ for every block of columns in order to fold in a single row. The older way to fold in one row, LINPACK’s dchud⁹ (Cholesky update), uses one Givens rotation per column, a $2 \times 2$ rotation that touches each column once, and would sit on the floor. LAPACK never picked it up, so from Python this is where it ends: the full single-row update, factor and right-hand side, is 2.0 ms against the QR form’s 2.5.

When the factor update wins

dtpqrt costs $np^2$, so it grows with the batch; refactorizing is a flat $p^3/3$ however many rows arrive.

For an online learner the factor update is the right default. Refactorizing only wins once a batch reaches a few hundred rows.

Where we ended up

At $p = 64$ the floor is the interpreter: a draw is 7 µs and the update is 34, almost all of it Python around a few microseconds of arithmetic, and no choice of routine does much more for it.

Say a decision has to come back in 25 ms. How many features can the model have? With the precision-form code I started with, dtrsm for the draw and numpy for the projection, about 6,000. With dtrsv and everything inside scipy, about 8,000:

A third more features in the same budget, from the same routines called the right way. And since the draw is now reading the factor at the speed main memory delivers it, that’s as many features as this budget buys on this hardware, which is the most useful thing to know.