2.3  Synthetic Regression Data

Before we can train a model we need data. Real datasets are what we ultimately care about, but they conflate three separate sources of failure: a misspecified model, a flawed optimization algorithm, and pathological data. When a method performs poorly on real data, these explanations can be difficult to distinguish. Synthetic data removes this ambiguity by construction. If we know the data-generating process exactly (the true weights \(\mathbf{w}^*\), the true bias \(b^*\), and the noise distribution), then any systematic failure to recover them (beyond the irreducible noise) points to the algorithm or implementation. This is why a compatible synthetic dataset is a useful early implementation test for a new learning method. We first confirm that the implementation solves a compatible problem with known parameters before evaluating it on real data.

%matplotlib inline
from d2l import torch as d2l
import torch
import random
%matplotlib inline
from d2l import tensorflow as d2l
import tensorflow as tf
import random
%matplotlib inline
from d2l import jax as d2l
import jax
from jax import numpy as jnp
import numpy as np
import random
import tensorflow as tf
%matplotlib inline
from d2l import mxnet as d2l
from mxnet import np, npx, gluon
import random
npx.set_np()

2.3.1 Generating the Dataset

For clarity, this example uses two-dimensional inputs. The following code snippet generates 1000 examples with 2-dimensional features drawn from a standard normal distribution. The resulting design matrix \(\mathbf{X}\) belongs to \(\mathbb{R}^{1000 \times 2}\). We generate each label by applying a ground truth linear function, corrupting them via additive noise \(\boldsymbol{\epsilon}\), drawn independently and identically for each example:

\[\mathbf{y}= \mathbf{X} \mathbf{w}^* + b^* + \boldsymbol{\epsilon}. \tag{2.3.1}\]

For convenience we assume that \(\boldsymbol{\epsilon}\) is drawn from a normal distribution with mean \(\mu= 0\) and standard deviation \(\sigma = 0.01\). We put the generation code in the __init__ method of a subclass of d2l.DataModule (introduced in Section 2.2.3), calling save_hyperparameters() so that every constructor argument (the parameters w and b, the noise level, the split sizes, and batch_size) is stored as an attribute and the dataset stays introspectable.

class SyntheticRegressionData(d2l.DataModule):
    """Synthetic data for linear regression."""
    def __init__(self, w, b, noise=0.01, num_train=1000, num_val=1000, 
                 batch_size=32):
        super().__init__()
        self.save_hyperparameters()
        n = num_train + num_val
        self.X = d2l.randn(n, len(w))
        eps = d2l.randn(n, 1) * noise
        self.y = d2l.matmul(self.X, d2l.reshape(w, (-1, 1))) + b + eps
class SyntheticRegressionData(d2l.DataModule):
    """Synthetic data for linear regression."""
    def __init__(self, w, b, noise=0.01, num_train=1000, num_val=1000, 
                 batch_size=32):
        super().__init__()
        self.save_hyperparameters()
        n = num_train + num_val
        self.X = tf.random.normal((n, w.shape[0]))
        eps = tf.random.normal((n, 1)) * noise
        self.y = d2l.matmul(self.X, d2l.reshape(w, (-1, 1))) + b + eps
class SyntheticRegressionData(d2l.DataModule):
    """Synthetic data for linear regression."""
    def __init__(self, w, b, noise=0.01, num_train=1000, num_val=1000,
                 batch_size=32, key=None):
        super().__init__()
        self.save_hyperparameters()
        # Resolve the key at call time rather than reusing a key in the signature.
        key = jax.random.key(0) if key is None else key
        n = num_train + num_val
        key1, key2 = jax.random.split(key)
        self.X = jax.random.normal(key1, (n, w.shape[0]))
        eps = jax.random.normal(key2, (n, 1)) * noise
        self.y = d2l.matmul(self.X, d2l.reshape(w, (-1, 1))) + b + eps
class SyntheticRegressionData(d2l.DataModule):
    """Synthetic data for linear regression."""
    def __init__(self, w, b, noise=0.01, num_train=1000, num_val=1000, 
                 batch_size=32):
        super().__init__()
        self.save_hyperparameters()
        n = num_train + num_val
        self.X = d2l.randn(n, len(w))
        eps = d2l.randn(n, 1) * noise
        self.y = d2l.matmul(self.X, d2l.reshape(w, (-1, 1))) + b + eps

Below, we set the true parameters to \(\mathbf{w}^* = [2, -3.4]^\top\) and \(b^* = 4.2\). Later, we can check our estimated parameters against these ground truth values.

data = SyntheticRegressionData(w=d2l.tensor([2, -3.4]), b=4.2)

Each row of data.X is a feature vector in \(\mathbb{R}^2\), and each row of data.y is a scalar label. We inspect the first entry.

print('features:', data.X[0],'\nlabel:', data.y[0])
features: tensor([-1.2309,  0.6609]) 
label: tensor([-0.5205])
features: tf.Tensor([-0.70784295 -0.2909985 ], shape=(2,), dtype=float32) 
label: tf.Tensor([3.7740877], shape=(1,), dtype=float32)
features: [ 1.0040143 -0.9063372] 
label: [9.265151]
features: [ 1.2680211 -0.9800729] 
label: [10.066712]

2.3.2 Reading the Dataset

Training machine learning models often requires multiple passes over a dataset, grabbing one minibatch of examples at a time. This data is then used to update the model. To illustrate how this works, we implement the get_dataloader method, registering it in the SyntheticRegressionData class via add_to_class (introduced in Section 2.2.1). It takes a batch size, a matrix of features, and a vector of labels, and generates minibatches of size batch_size. As such, each minibatch consists of a tuple of features and labels. Note that we need to be mindful of whether we’re in training or validation mode: in the former, we will want to read the data in random order, whereas for the latter, being able to read data in a pre-defined order may be important for debugging purposes.

@d2l.add_to_class(SyntheticRegressionData)
def get_dataloader(self, train):
    if train:
        indices = list(range(0, self.num_train))
        # The examples are read in random order
        random.shuffle(indices)
    else:
        indices = list(range(self.num_train, self.num_train+self.num_val))
    for i in range(0, len(indices), self.batch_size):
        batch_indices = d2l.tensor(indices[i: i+self.batch_size])
        yield self.X[batch_indices], self.y[batch_indices]
@d2l.add_to_class(SyntheticRegressionData)
def get_dataloader(self, train):
    if train:
        indices = list(range(0, self.num_train))
        # The examples are read in random order
        random.shuffle(indices)
    else:
        indices = list(range(self.num_train, self.num_train+self.num_val))
    for i in range(0, len(indices), self.batch_size):
        j = tf.constant(indices[i : i+self.batch_size])
        yield tf.gather(self.X, j), tf.gather(self.y, j)
@d2l.add_to_class(SyntheticRegressionData)
def get_dataloader(self, train):
    if train:
        indices = list(range(0, self.num_train))
        # The examples are read in random order
        random.shuffle(indices)
    else:
        indices = list(range(self.num_train, self.num_train+self.num_val))
    for i in range(0, len(indices), self.batch_size):
        batch_indices = d2l.tensor(indices[i: i+self.batch_size])
        yield self.X[batch_indices], self.y[batch_indices]
@d2l.add_to_class(SyntheticRegressionData)
def get_dataloader(self, train):
    if train:
        indices = list(range(0, self.num_train))
        # The examples are read in random order
        random.shuffle(indices)
    else:
        indices = list(range(self.num_train, self.num_train+self.num_val))
    for i in range(0, len(indices), self.batch_size):
        batch_indices = d2l.tensor(indices[i: i+self.batch_size])
        yield self.X[batch_indices], self.y[batch_indices]

We inspect the first minibatch. The feature shape records the minibatch size and input dimensionality, while the label shape has the same leading dimension.

X, y = next(iter(data.train_dataloader()))
print('X shape:', X.shape, '\ny shape:', y.shape)
X shape: torch.Size([32, 2]) 
y shape: torch.Size([32, 1])
X shape: (32, 2) 
y shape: (32, 1)
X shape: (32, 2) 
y shape: (32, 1)
X shape: (32, 2) 
y shape: (32, 1)

Iterating over data.train_dataloader() yields distinct minibatches until the dataset is exhausted (try it). Writing the loader by hand makes every step explicit, but it costs us in three ways: all of the data must fit in memory, the iteration is single-threaded Python looping over indices, and there is no prefetching to overlap data loading with computation on the previous batch. The data loaders built into a deep learning framework fix all three. They run several worker processes in parallel, prefetch the next batch while the current one trains, and stream from sources such as files, network streams, or generators that produce data on the fly. We now switch to the framework’s built-in loader, which presents an identical interface to the caller.

2.3.3 Concise Implementation of the Data Loader

Rather than writing our own iterator, we can call the existing API in a framework to load data. As before, we need a dataset with features X and labels y. Beyond that, we set batch_size in the built-in data loader and let it take care of shuffling examples efficiently.

JAX offers a NumPy-like API with device acceleration and functional transformations, and at least the current version ships no data loading methods of its own. Other libraries already provide great data loaders, and JAX suggests using them instead. Here we will grab TensorFlow’s data loader and modify it slightly to make it work with JAX.

@d2l.add_to_class(d2l.DataModule)
def get_tensorloader(self, tensors, train, indices=slice(0, None)):
    tensors = tuple(a[indices] for a in tensors)
    dataset = torch.utils.data.TensorDataset(*tensors)
    return torch.utils.data.DataLoader(dataset, self.batch_size,
                                       shuffle=train)
@d2l.add_to_class(d2l.DataModule)
def get_tensorloader(self, tensors, train, indices=slice(0, None)):
    tensors = tuple(a[indices] for a in tensors)
    shuffle_buffer = tensors[0].shape[0] if train else 1
    return tf.data.Dataset.from_tensor_slices(tensors).shuffle(
        buffer_size=shuffle_buffer).batch(self.batch_size)
class TensorFlowDataLoader:
    """Expose a tf.data.Dataset as re-iterable NumPy batches."""
    def __init__(self, dataset):
        self.dataset = dataset

    def __iter__(self):
        return self.dataset.as_numpy_iterator()

    def __len__(self):
        return len(self.dataset)

@d2l.add_to_class(d2l.DataModule)
def get_tensorloader(self, tensors, train, indices=slice(0, None)):
    tensors = tuple(a[indices] for a in tensors)
    # Use TensorFlow's data loader. JAX and Flax do not provide data-loading
    # functionality. `drop_remainder=train` keeps every
    # *training* minibatch the same shape, so a `@jax.jit`'d step
    # function compiles once per epoch instead of recompiling for the
    # smaller last batch.
    shuffle_buffer = tensors[0].shape[0] if train else 1
    dataset = tf.data.Dataset.from_tensor_slices(tensors).shuffle(
        buffer_size=shuffle_buffer).batch(
            self.batch_size, drop_remainder=train)
    return TensorFlowDataLoader(dataset)
@d2l.add_to_class(d2l.DataModule)
def get_tensorloader(self, tensors, train, indices=slice(0, None)):
    tensors = tuple(a[indices] for a in tensors)
    dataset = gluon.data.ArrayDataset(*tensors)
    return gluon.data.DataLoader(dataset, self.batch_size,
                                 shuffle=train)
@d2l.add_to_class(SyntheticRegressionData)
def get_dataloader(self, train):
    i = slice(0, self.num_train) if train else slice(self.num_train, None)
    return self.get_tensorloader((self.X, self.y), train, i)

The new data loader behaves just like the previous one, except that it is more efficient and has some added functionality.

X, y = next(iter(data.train_dataloader()))
print('X shape:', X.shape, '\ny shape:', y.shape)
X shape: torch.Size([32, 2]) 
y shape: torch.Size([32, 1])
X shape: (32, 2) 
y shape: (32, 1)
X shape: (32, 2) 
y shape: (32, 1)
X shape: (32, 2) 
y shape: (32, 1)

For instance, the data loader provided by the framework API supports the built-in __len__ method, so we can query its length, i.e., the number of batches.

len(data.train_dataloader())
32
32
31
32

With 1000 training examples and a batch size of 32, we expect \(\lceil 1000 / 32 \rceil = 32\) batches: 31 full ones and a final partial batch of 8 examples. Note also that the built-in training loader reshuffles the examples at the start of every epoch, just as our hand-rolled loader drew a fresh random order on each call; exercise 9 of Section 2.4 asks why this reshuffling matters.

The JAX loader reports 31 batches rather than 32. This is because get_tensorloader passes drop_remainder=True when training: the final partial batch of 8 examples is discarded. We do this so that every training minibatch has an identical shape, which keeps a @jax.jit-compiled training step from being recompiled for the differently sized last batch. The price is that we drop a handful of examples each epoch, which is negligible here. A loader that keeps the partial batch would report 32.

2.3.4 Summary

Synthetic data lets us check the recovered parameters against the truth we fixed: because we chose \(\mathbf{w}^*\) and \(b^*\) ourselves, we can see after training whether the estimates agree, which makes such datasets an early place to validate an algorithm against a compatible known-answer problem. Passing this check does not establish robustness to misspecification or real deployment data. The SyntheticRegressionData class introduced here packages this data-generating process as a DataModule subclass, separating where the batches come from from how a model consumes them. Along the way we implemented the same get_dataloader protocol twice: a transparent hand-rolled iterator that is easy to read but loads everything in memory and loops in Python, and a framework-native loader that shuffles, prefetches, and parallelizes for us. The hand-rolled version is there to teach; the framework version is what we use from here on.

2.3.5 Exercises

  1. Partial batches. With 1000 examples per split and batch_size=32, each loader ends a pass with a partial batch of 8 examples. DataLoader discards that batch when given drop_last=True.
    1. The training loader reshuffles every epoch, whereas the validation loader reads in a fixed order. For each loader, state which examples a discarded partial batch leaves out over many epochs, and decide for which loader discarding is acceptable.
    2. Suppose that the partial batch is kept and that the validation loss is reported as the unweighted mean of the 32 per-batch mean losses. Compute the weight that this mean gives one example of the partial batch and one example of a full batch, and compare both with the weight \(1/1000\) in the exact mean over the validation set.
  1. Partial batches. With 1000 examples per split and batch_size=32, each loader ends a pass with a partial batch of 8 examples. Dataset.batch discards that batch when given drop_remainder=True.
    1. The training loader reshuffles every epoch, whereas the validation loader reads in a fixed order. For each loader, state which examples a discarded partial batch leaves out over many epochs, and decide for which loader discarding is acceptable.
    2. Suppose that the partial batch is kept and that the validation loss is reported as the unweighted mean of the 32 per-batch mean losses. Compute the weight that this mean gives one example of the partial batch and one example of a full batch, and compare both with the weight \(1/1000\) in the exact mean over the validation set.
  1. Partial batches. With 1000 examples per split and batch_size=32, each loader ends a pass with a partial batch of 8 examples. get_tensorloader passes drop_remainder=train, so the training loader discards that batch while the validation loader keeps it.
    1. The training loader reshuffles every epoch, whereas the validation loader reads in a fixed order. State which training examples the discarded batch leaves out over many epochs, and which examples every validation pass would leave out if drop_remainder=True were passed for validation as well. Explain why discarding is acceptable for one loader and not for the other.
    2. Suppose that the validation loss is reported as the unweighted mean of the 32 per-batch mean losses. Compute the weight that this mean gives one example of the partial batch and one example of a full batch, and compare both with the weight \(1/1000\) in the exact mean over the validation set.
  1. Partial batches. With 1000 examples per split and batch_size=32, each loader ends a pass with a partial batch of 8 examples. Gluon’s DataLoader discards that batch when given last_batch='discard'.
    1. The training loader reshuffles every epoch, whereas the validation loader reads in a fixed order. For each loader, state which examples a discarded partial batch leaves out over many epochs, and decide for which loader discarding is acceptable.
    2. Suppose that the partial batch is kept and that the validation loss is reported as the unweighted mean of the 32 per-batch mean losses. Compute the weight that this mean gives one example of the partial batch and one example of a full batch, and compare both with the weight \(1/1000\) in the exact mean over the validation set.
  1. Data beyond memory. ● Suppose that len(w) and num_train + num_val are so large that \(\mathbf{X}\) no longer fits in memory, so __init__ generates it in chunks and writes it to disk.
    1. Name the statements in SyntheticRegressionData.__init__ and in the hand-rolled get_dataloader whose memory or cost grows with the number of examples, and state which of them writing \(\mathbf{X}\) to disk in chunks does not fix.
    2. Reading single examples from random disk locations is slow, so the training loader should read long contiguous blocks. Design a loader that reads such blocks and still visits the examples in a fresh order every epoch. State in what sense this order falls short of a uniformly random permutation.
    3. The loader must also avoid storing a permutation table with one entry per example. Explain how a pseudorandom permutation (Naor and Reingold 1999) meets this requirement while still changing the order every epoch.
  2. [code] On-the-fly generation. Write a DataModule subclass that stores no \(\mathbf{X}\) or \(\mathbf{y}\). Its training loader instead draws every minibatch fresh from the data-generating process when the minibatch is requested.
    1. A stream of fresh minibatches has no natural end, yet Trainer.prepare_data in Section 2.2.4 calls len on both loaders. Decide how many minibatches one epoch yields, and make len(data.train_dataloader()) return that number.
    2. Decide whether the validation loader should also draw fresh examples on every call, and implement your choice. Justify it by what the validation loss is used for.
  1. Reproducibility. Two instances of SyntheticRegressionData built with the same arguments hold different X and y. Make them identical by fixing the seed with torch.manual_seed. State where the call must be placed relative to the draws of X and eps, and explain why seeding once at the top of a notebook does not make the second instance equal to the first.
  1. Reproducibility. Two instances of SyntheticRegressionData built with the same arguments hold different X and y. Make them identical by fixing the seed with tf.random.set_seed. State where the call must be placed relative to the draws of X and eps, and explain why seeding once at the top of a notebook does not make the second instance equal to the first.
  1. Reproducibility. SyntheticRegressionData splits its key into key1 for X and key2 for eps.
    1. Explain why two instances built with the same key hold identical data without any global seed.
    2. Change both draws to use key itself instead of key1 and key2, and compare eps / noise with the entries of X numerically. State what the split guarantees about X and eps that this version does not.
    3. Construct five instances whose datasets are independent of each other, as averaging an error over repeated trials requires, and state how you derive their keys.
  1. Reproducibility. Two instances of SyntheticRegressionData built with the same arguments hold different X and y. Make them identical by fixing the seed with npx.random.seed. State where the call must be placed relative to the draws of X and eps, and explain why seeding once at the top of a notebook does not make the second instance equal to the first.
  1. [code] Recovery under noise. On the training split of SyntheticRegressionData, fit \(\hat{\mathbf{w}}\) and \(\hat{b}\) jointly by least squares, appending a column of ones to \(\mathbf{X}\) as in Section 2.1.1.3 and solving with the normal equations or lstsq. Measure the distance \(\|\hat{\mathbf{w}} - \mathbf{w}^*\|_2\) from the generating weights \(\mathbf{w}^*\), averaged over 10 independent datasets for each setting.
    1. Before running, state how you expect this error to scale with the noise standard deviation noise and with num_train.
    2. Vary noise over \(\{0.001, 0.01, 0.1, 0.5, 1.0\}\) with num_train=1000, and plot the error against the noise level on logarithmic axes.
    3. Vary num_train over \(5, 10, 20, 50, \ldots, 10{,}000\) with noise=0.01, and plot the error against num_train on logarithmic axes.
  2. Planted bug. A faulty variant of the hand-rolled get_dataloader shuffles the rows of X and the entries of y with two independent permutations. We collect the pairs that it yields in one training epoch on data, whose parameters are \(\mathbf{w}^* = [2, -3.4]^\top\) and \(b^* = 4.2\) with noise standard deviation \(0.01\). We then fit \(\hat{\mathbf{w}}\) and \(\hat{b}\) to these pairs by least squares.
    1. Predict the approximate values of \(\hat{\mathbf{w}}\) and \(\hat{b}\), and explain your prediction.
    2. Estimate the mean squared error of this fit on the collected pairs, and compare it with the smallest mean squared error that any predictor can attain on correctly paired data.