Show HN: I implemented a neural network in SQL

Show HN: I implemented a neural network in SQL

The network architecture is 196 -> 32 tanh -> 10 softmax, with N_SAMPLES = 700, TRAIN_FRAC = 0.7, LR = 0.5, STEPS = 60, and CHUNK = 250.

Zero pixel optimization

Drop zero-valued pixels from the (dominant) layer-0 contraction. A background pixel contributes 0 * weight = 0, so skipping those rows shrinks the join exactly - the result is identical, and the speedup scales with the fraction of zeros (a dark background). On dense inputs it is a no-op.

Measured ~1.8x on real Fashion-MNIST (~50% zero pixels): 2.56 -> 1.45 s/step. SKIP_ZERO_PIXELS = True.

Data loading

def fashion_mnist():
    """The whole training set, left lazy so SQL streams and samples it.
    The real path returns a dask-backed (chunked) Dataset - nothing is pulled
    into memory here; ``from_dataset`` reads it chunk by chunk on demand, and
    the random subsample happens later in SQL. The offline fallback is a small
    synthetic set built in memory.
    """
    try:
        ds = xr.open_dataset(
            "s3://carbonplan-share/xbatcher/fashion-mnist-train.zarr",
            engine="zarr",
            chunks=None,
            backend_kwargs={"storage_options": {"anon": True}},
        )
        if "channel" in ds.dims:
            ds = ds.isel(channel=0, drop=True)
        # To float64, lazily (no full read). This zarr already stores images
        # as float in [0, 1]; only integer-encoded sources ([0, 255]) rescale.
        images = ds["images"].astype("float64")
        if not np.issubdtype(ds["images"].dtype, np.floating):
            images = images / 255.0
        ds = ds.assign(images=images, labels=ds["labels"].astype("int64"))
    except Exception:
        # Offline fallback: a separable synthetic set (per-class template +
        # noise), so the same pipeline still learns without the network. A pool
        # larger than N_SAMPLES so the SQL subsample still has something to pick.
        rng = np.random.default_rng(0)
        n = 3 * N_SAMPLES
        templates = rng.standard_normal((10, SIDE, SIDE))
        labels = rng.integers(0, 10, n).astype("int64")
        images = templates[labels] + 0.6 * rng.standard_normal((n, SIDE, SIDE))
        ds = xr.Dataset(
            {
                "images": (("sample", "height", "width"), images),
                "labels": (("sample",), labels),
            }
        )
    # Integer index coords are the SQL join keys (sample, height, width).
    return ds[["images", "labels"]].assign_coords(
        sample=np.arange(ds.sizes["sample"]),
        height=np.arange(ds.sizes["height"]),
        width=np.arange(ds.sizes["width"]),
    )

Model construction

def build_model_with_table_names(
    init_weight: Callable[[int, int], np.ndarray],
    init_bias: Callable[[int], np.ndarray],
    widths=WIDTHS,
) -> tuple[xr.Dataset, dict[tuple[str, ...], str]]:
    """The network as one Dataset that splits into tables per layer.
    Layer ``i`` is a weight matrix ``layer_i (inp_i, out_i)`` and a separate
    bias vector ``bias_i (out_i,)``.
    """
    weights = {
        f"layer_{i}": ((f"inp_{i}", f"out_{i}"), init_weight(inp, out))
        for i, (inp, out) in enumerate(zip(widths[:-1], widths[1:]))
    }
    biases = {
        f"bias_{i}": ((f"out_{i}",), init_bias(out))
        for i, out in enumerate(widths[1:])
    }
    coords = {}
    coords.update(
        {f"inp_{i}": np.arange(inp) for i, inp in enumerate(widths[:-1])}
    )
    coords.update(
        {f"out_{i}": np.arange(out) for i, out in enumerate(widths[1:])}
    )
    ds = xr.Dataset({**weights, **biases}, coords=coords)
    names: dict[tuple[str, ...], str] = {}
    for i in range(len(weights)):
        names[(f"inp_{i}", f"out_{i}")] = f"layer{i}"
        names[(f"out_{i}",)] = f"bias{i}"
    return ds, names

Main training loop

def main():
    rng = np.random.default_rng(1)
    mnist = fashion_mnist()
    ctx = xql.XarrayContext()

    # One Dataset splits into two tables: pixels (sample, height, width) and
    # labels (sample). The dim names are the join keys.
    ctx.from_dataset(
        "mnist",
        mnist,
        chunks=dict(sample=CHUNK),
        table_names={
            ("sample", "height", "width"): "pixels",
            ("sample",): "labels",
        },
    )

    # Draw a random N_SAMPLES subset in SQL (ORDER BY random() LIMIT), carrying
    # each sample's label and a train/test tag. `data` is the working label
    # table: cache() pins the chosen subset so every downstream query sees the
    # same split without rescanning the source. `ORDER BY random()` shuffles the
    # whole label column, so the subset is order-independent even if the on-disk
    # samples are class-sorted.
    data = ctx.sql(f"""
        SELECT sample, labels,
               CASE WHEN random() < {TRAIN_FRAC} THEN 'train' ELSE 'test' END AS split
        FROM labels
        ORDER BY random()
        LIMIT {N_SAMPLES}
    """).cache()
    ctx.deregister_table("data")
    ctx.register_table("data", data)

The forward pass runs over ALL samples: train rows drive learning, test rows ride along so we can score them from the same logits. Only delta2 is restricted to train, so the gradients (and the trained weights) are identical to a train-only forward - test is never backpropagated.

    fwd0 = ctx.sql(f"""
        WITH c AS (
            -- z = x @ W: matmul of the input and first weight matrix
            SELECT a.sample, w.out AS out, SUM(a.val * w.val) AS z
            FROM (
                SELECT sample, height * {SIDE} + width AS inp, images AS val
                FROM pixels {zero_where}
            ) a
            JOIN weight w ON a.inp = w.inp AND w.layer = 0
            GROUP BY a.sample, w.out
        )
        -- activation(z + b): Add in the bias term, then perform activation
        SELECT c.sample, c.out AS out, c.z + b.val AS z, tanh(c.z + b.val) AS val
        FROM c
        JOIN bias b ON c.out = b.out AND b.layer = 0
    """).cache()

The backward pass computes output error delta2 = softmax(logits) - onehot(label), then propagates through each layer using the chain rule. Weight and bias gradients are computed as averages over the training set.

    delta2 = ctx.sql(f"""
        WITH m AS (SELECT sample, MAX(z) AS m FROM logits GROUP BY sample),
             e AS (SELECT logits.sample, logits.out, exp(logits.z - m.m) AS e
                   FROM logits JOIN m ON logits.sample = m.sample),
             s AS (SELECT sample, SUM(e) AS s FROM e GROUP BY sample)
        SELECT e.sample, e.out, e.e / s.s - CASE WHEN e.out = y.labels THEN 1.0 ELSE 0.0 END AS val
        FROM e
        JOIN s ON e.sample = s.sample
        JOIN data y ON y.sample = e.sample
        -- restrict the error to train, so every downstream gradient is train-only
        WHERE e.sample IN (SELECT sample FROM data WHERE split = 'train')
    """).cache()

SGD updates are applied in a single query per relation:

    ctx.sql(f"""
        UPDATE weight SET val = val - {LR} * g.val
        FROM g0 g
        WHERE weight.inp = g.inp AND weight.out = g.out AND weight.layer = 0
    """)

Accuracy tracking

    acc = ctx.sql(f"""
        WITH m AS (SELECT sample, MAX(z) AS m FROM logits GROUP BY sample),
             p AS (SELECT logits.sample, logits.out AS pred
                   FROM logits JOIN m ON logits.sample = m.sample
                   WHERE logits.z = m.m)
        SELECT p.sample, p.pred, d.labels, d.split
        FROM p JOIN data d ON p.sample = d.sample
    """)

Final output

    print(f"Training complete - {STEPS} steps, final accuracy {acc_train*100:.1f}% / {acc_test*100:.1f}%")
    print(f"Trained weights: {dict(trained.sizes)}.")
    print(trained)
    trained.to_zarr(
        f"fashion_mnist_mlp_"
        f"{datetime.datetime.now().isoformat(timespec='seconds')}.zarr"
    )