Distributed Linear and Logistic Regression

"They asked each of us to summarize our shard in a single small table. We added the tables together, solved once, and the answer was exactly right. I have never felt so useful and so replaceable at the same time."

A Worker That Emitted Its Sufficient Statistics

In the previous chapters of Part III we built the general machinery of distributed training: gradient aggregation and its convergence behavior in Chapter 10, and the parameter-server pattern for sharded model state in Chapter 11. Now we apply that machinery to specific models, starting with the two oldest and most transparent ones. Linear and logistic regression are not interesting because they are powerful; they are interesting because they expose, with nothing hidden, why some models distribute perfectly in a single communication round while others require many. The dividing line is whether the training computation collapses into a sum of per-example statistics, and once you can see that line you can predict the distributed cost of a method before writing a line of code.

1. Linear Regression as a Sum Over Examples Beginner

Ordinary least squares fits a weight vector $w$ that minimizes the squared error between predictions $Xw$ and targets $y$, where $X$ is the $N \times d$ design matrix of $N$ examples and $d$ features. Setting the gradient of the squared-error objective to zero gives the normal equations, a linear system whose solution is the global optimum in closed form:

Stare at the two quantities the solution actually needs. The matrix $X^\top X$ is $d \times d$ and the vector $X^\top y$ is $d \times 1$; neither grows with the number of examples $N$. More importantly, both are sums over examples. Writing $x_i$ for the $i$-th row of $X$ and $y_i$ for its target,

This is the whole game. A sum over examples does not care how the examples are grouped, so we can split the $N$ rows into $K$ disjoint shards, hand shard $k$ to worker $k$, and let each worker accumulate its own partial sums $\sum_{i \in S_k} x_i x_i^\top$ and $\sum_{i \in S_k} x_i y_i$ over only the data it holds. Each worker then ships a tiny fixed-size pair of statistics, $d^2 + d$ numbers, regardless of how many millions of rows its shard contained. Adding the $K$ partial matrices and the $K$ partial vectors reconstructs the exact global $X^\top X$ and $X^\top y$, and a single solve yields the same $w^\star$ a single machine would have computed. The combine step is summing one array per worker, which is precisely the all-reduce primitive of Section 4.3, and the per-shard accumulate-then-merge structure is the mergeable aggregation pattern of Section 6.8 applied to matrices instead of counts.

2. The Per-Shard Statistics Picture Beginner

Figure 12.1.1 makes the dataflow concrete. The full dataset is partitioned by rows across workers; each worker reduces its shard, no matter how large, to one small $X^\top X$ block and one short $X^\top y$ vector; a single all-reduce sums those blocks into the global statistics; and one solve on the summed system produces the exact model. The picture is worth holding in mind because nearly every method in this chapter is a variation on it, differing only in what the per-shard statistic is and whether one round suffices.

3. The Exact One-Round Algorithm, From Scratch Intermediate

The code below realizes Figure 12.1.1 in pure NumPy. It builds a regression dataset, computes the single-machine normal-equations solution as ground truth, then simulates $K$ workers that each see only their shard, emit the pair $\left(X^\top X_k,\, X^\top y_k\right)$, sum those statistics (the all-reduce), and solve once. It also reports the communication accounting: how many numbers each worker sends versus how many it would send if it naively shipped its raw rows instead.

import numpy as np

rng = np.random.default_rng(7)
N, d, K = 200_000, 12, 16          # examples, features (incl. bias), workers
X = rng.standard_normal((N, d))
X[:, 0] = 1.0                       # intercept column
w_true = rng.standard_normal(d)
y = X @ w_true + 0.5 * rng.standard_normal(N)

# --- Single-machine normal equations: w = (X^T X)^{-1} X^T y ---
XtX_full = X.T @ X
Xty_full = X.T @ y
w_single = np.linalg.solve(XtX_full, Xty_full)

# --- Distributed: each worker sees ONLY its shard, emits sufficient statistics ---
shards = np.array_split(np.arange(N), K)
partials = []
for s in shards:
    Xs, ys = X[s], y[s]
    partials.append((Xs.T @ Xs, Xs.T @ ys))      # (d x d), (d,) per worker

# all-reduce: sum the per-shard statistics across workers (one round)
XtX_dist = sum(p[0] for p in partials)
Xty_dist = sum(p[1] for p in partials)
w_dist = np.linalg.solve(XtX_dist, Xty_dist)

# communication accounting
floats_per_worker = d * d + d
print(f"examples N             : {N}")
print(f"features d             : {d}")
print(f"workers K              : {K}")
print(f"floats sent / worker   : {floats_per_worker}  (d*d + d, independent of N)")
print(f"naive raw-data floats  : {N * d}  (if shards shipped data instead)")
print(f"max|w_dist - w_single| : {np.max(np.abs(w_dist - w_single)):.3e}")
print(f"relative error         : {np.linalg.norm(w_dist - w_single)/np.linalg.norm(w_single):.3e}")
print(f"identical solution     : {np.allclose(w_dist, w_single)}")

examples N             : 200000
features d             : 12
workers K              : 16
floats sent / worker   : 156  (d*d + d, independent of N)
naive raw-data floats  : 2400000  (if shards shipped data instead)
max|w_dist - w_single| : 5.818e-14
relative error         : 1.825e-14
identical solution     : True

The result in Output 12.1.1 is the same kind of exactness we met for the data-parallel gradient in Section 1.3, now for an entire closed-form model rather than a single gradient step. The model is identical, the communication is one round, and the bytes per worker are fixed at 156 numbers no matter how many rows the shard held. This is the gold standard a distributed learner can aspire to, and it is achievable precisely because least squares is a sum over examples with a closed-form solve. The relevant cost is no longer accuracy; it is the $O(d^2)$ size of the statistic and the $O(d^3)$ solve, which is why this exact route is ideal for many features in moderate dimension and gives way to iteration when $d$ is very large.

4. Logistic Regression Has No Postcard, But It Has a Gradient Intermediate

Logistic regression replaces the squared-error objective with the log-loss of a sigmoid, $\sigma(z) = 1/(1 + e^{-z})$, which models the probability of a binary label. Its objective and gradient are

There is no closed-form solver here: the sigmoid makes the normal-equations trick unavailable, so we cannot reduce the whole fit to one fixed table and one solve. But look at the gradient. It is, once again, a sum over examples. At any current weight $w$, each worker can compute its partial gradient over only its shard, the all-reduce sums those partials, and one synchronized gradient-descent step moves $w$. The catch is that the statistic now depends on the current $w$, so it is not sufficient once and for all; it must be recomputed and re-communicated at every iteration. Distributed logistic regression is therefore exact per step but iterative overall, which is exactly the distributed-SGD setting developed in Section 10.3. Each round is an all-reduce of a length-$d$ gradient; convergence takes many rounds rather than one.

When $d$ is modest, you can do better than plain gradient descent by aggregating richer per-shard statistics. Distributed L-BFGS keeps the communication pattern identical, an all-reduce of the summed gradient, but uses a limited-memory curvature approximation built from a short history of those aggregated gradients to take far better-conditioned steps, converging in tens of rounds where SGD might need thousands. The lesson generalizes cleanly: a model with a closed-form sufficient statistic distributes in one round; a model whose only summable quantity is the gradient distributes in many rounds, one all-reduce each. The number of rounds, not the existence of a combine step, is what separates these two regression cousins.

5. The General Lesson: Summable Statistics Are Cheap to Distribute Intermediate

The two regressions in this section bracket a principle that orders the entire chapter. If a model's training reduces to summing a fixed-size per-example statistic, and that summed statistic determines the model in closed form, distribution is exact and costs one all-reduce; linear regression is the archetype. If the only summable quantity is a gradient that depends on the current parameters, distribution is still exact per step but iterative, costing one all-reduce per round; logistic regression is the archetype. The communication payload in both cases is bounded by the statistic's size, not the data's, which is why classical models with low-dimensional sufficient statistics scale so gracefully even on enormous datasets.

This framing tells you what to look for in every later method. Naive Bayes and Gaussian discriminant models reduce to per-class sums and counts, one round like linear regression. The mean update in k-means clustering (Section 12.6) is a sum of points per cluster, exact per iteration. Tree-based methods aggregate histograms of feature-value statistics rather than raw rows, the trick that makes distributed gradient boosting practical in Section 12.5. In each case the question is identical to the one we asked here: what fixed-size quantity does a worker contribute to the shared sum, and does one sum finish the job or must we iterate. Answering it predicts the distributed cost before any code is written.

from pyspark.ml.regression import LinearRegression
from pyspark.ml.classification import LogisticRegression

# `train` is a distributed Spark DataFrame with columns "features", "label".
lin = LinearRegression(featuresCol="features", labelCol="label").fit(train)   # normal-equations / IRLS, distributed
log = LogisticRegression(featuresCol="features", labelCol="label",
                         maxIter=100).fit(train)                              # distributed L-BFGS under the hood

Linear and logistic regression, the humblest models in the book, have handed us the organizing idea for distributed classical machine learning: find the summable statistic, count the rounds. The next section carries the same question to support vector machines, where the per-example contributions are constrained by a quadratic program and the combine step becomes more interesting. That tour begins in Section 12.2.