ddpol.py

   1 #!/usr/bin/env python
   2
   3 # Any copyright is dedicated to the Public Domain.
   4 # https://creativecommons.org/publicdomain/zero/1.0/
   5
   6 # Written by Francois Fleuret <francois@fleuret.org>
   7
   8 import math, argparse
   9 import matplotlib.pyplot as plt
  10
  11 import torch
  12
  13 ######################################################################
  14
  15 parser = argparse.ArgumentParser(description='Example of double descent with polynomial regression.')
  16
  17 parser.add_argument('--D-max',
  18                     type = int, default = 16)
  19
  20 parser.add_argument('--nb-runs',
  21                     type = int, default = 250)
  22
  23 parser.add_argument('--nb-train-samples',
  24                     type = int, default = 8)
  25
  26 parser.add_argument('--train-noise-std',
  27                     type = float, default = 0.)
  28
  29 parser.add_argument('--seed',
  30                     type = int, default = 0,
  31                     help = 'Random seed (default 0, < 0 is no seeding)')
  32
  33 args = parser.parse_args()
  34
  35 if args.seed >= 0:
  36     torch.manual_seed(args.seed)
  37
  38 ######################################################################
  39
  40 def pol_value(alpha, x):
  41     x_pow = x.view(-1, 1) ** torch.arange(alpha.size(0)).view(1, -1)
  42     return x_pow @ alpha
  43
  44 def fit_alpha(x, y, D, a = 0, b = 1, rho = 1e-12):
  45     M = x.view(-1, 1) ** torch.arange(D + 1).view(1, -1)
  46     B = y
  47
  48     if D >= 2:
  49         q = torch.arange(2, D + 1, dtype = x.dtype).view(1, -1)
  50         r = q.view(-1,  1)
  51         beta = x.new_zeros(D + 1, D + 1)
  52         beta[2:, 2:] = (q-1) * q * (r-1) * r * (b**(q+r-3) - a**(q+r-3))/(q+r-3)
  53         l, U = beta.eig(eigenvectors = True)
  54         Q = U @ torch.diag(l[:, 0].clamp(min = 0) ** 0.5) # clamp deals with ~0 negative values
  55         B = torch.cat((B, y.new_zeros(Q.size(0))), 0)
  56         M = torch.cat((M, math.sqrt(rho) * Q.t()), 0)
  57
  58     return torch.lstsq(B, M).solution[:D+1, 0]
  59
  60 ######################################################################
  61
  62 # The "ground truth"
  63
  64 def phi(x):
  65     return torch.abs(torch.abs(x - 0.4) - 0.2) + x/2 - 0.1
  66
  67 ######################################################################
  68
  69 def compute_mse(nb_train_samples):
  70     mse_train = torch.zeros(args.nb_runs, args.D_max + 1)
  71     mse_test = torch.zeros(args.nb_runs, args.D_max + 1)
  72
  73     for k in range(args.nb_runs):
  74         x_train = torch.rand(nb_train_samples, dtype = torch.float64)
  75         y_train = phi(x_train)
  76         if args.train_noise_std > 0:
  77             y_train = y_train + torch.empty_like(y_train).normal_(0, args.train_noise_std)
  78         x_test = torch.linspace(0, 1, 100, dtype = x_train.dtype)
  79         y_test = phi(x_test)
  80
  81         for D in range(args.D_max + 1):
  82             alpha = fit_alpha(x_train, y_train, D)
  83             mse_train[k, D] = ((pol_value(alpha, x_train) - y_train)**2).mean()
  84             mse_test[k, D] = ((pol_value(alpha, x_test) - y_test)**2).mean()
  85
  86     return mse_train.median(0).values, mse_test.median(0).values
  87
  88 ######################################################################
  89 # Plot the MSE vs. degree curves
  90
  91 fig = plt.figure()
  92
  93 ax = fig.add_subplot(1, 1, 1)
  94 ax.set_yscale('log')
  95 ax.set_ylim(1e-5, 1)
  96 ax.set_xlabel('Polynomial degree', labelpad = 10)
  97 ax.set_ylabel('MSE', labelpad = 10)
  98
  99 ax.axvline(x = args.nb_train_samples - 1,
 100            color = 'gray', linewidth = 0.5, linestyle = '--')
 101
 102 ax.text(args.nb_train_samples - 1.2, 1e-4, 'Nb. params = nb. samples',
 103         fontsize = 10, color = 'gray',
 104         rotation = 90, rotation_mode='anchor')
 105
 106 mse_train, mse_test = compute_mse(args.nb_train_samples)
 107 ax.plot(torch.arange(args.D_max + 1), mse_train, color = 'blue', label = 'Train')
 108 ax.plot(torch.arange(args.D_max + 1), mse_test, color = 'red', label = 'Test')
 109
 110 ax.legend(frameon = False)
 111
 112 fig.savefig('dd-mse.pdf', bbox_inches='tight')
 113
 114 plt.close(fig)
 115
 116 ######################################################################
 117 # Plot multiple MSE vs. degree curves
 118
 119 fig = plt.figure()
 120
 121 ax = fig.add_subplot(1, 1, 1)
 122 ax.set_yscale('log')
 123 ax.set_ylim(1e-5, 1)
 124 ax.set_xlabel('Polynomial degree', labelpad = 10)
 125 ax.set_ylabel('MSE', labelpad = 10)
 126
 127 nb_train_samples_min = args.nb_train_samples - 4
 128 nb_train_samples_max = args.nb_train_samples
 129
 130 for nb_train_samples in range(nb_train_samples_min, nb_train_samples_max + 1, 2):
 131     mse_train, mse_test = compute_mse(nb_train_samples)
 132     e = float(nb_train_samples - nb_train_samples_min) / float(nb_train_samples_max - nb_train_samples_min)
 133     e = 0.15 + 0.7 * e
 134     ax.plot(torch.arange(args.D_max + 1), mse_train, color = (e, e, 1.0), label = f'Train N={nb_train_samples}')
 135     ax.plot(torch.arange(args.D_max + 1), mse_test, color = (1.0, e, e), label = f'Test N={nb_train_samples}')
 136
 137 ax.legend(frameon = False)
 138
 139 fig.savefig('dd-multi-mse.pdf', bbox_inches='tight')
 140
 141 plt.close(fig)
 142
 143 ######################################################################
 144 # Plot some examples of train / test
 145
 146 torch.manual_seed(9) # I picked that for pretty
 147
 148 x_train = torch.rand(args.nb_train_samples, dtype = torch.float64)
 149 y_train = phi(x_train)
 150 if args.train_noise_std > 0:
 151     y_train = y_train + torch.empty_like(y_train).normal_(0, args.train_noise_std)
 152 x_test = torch.linspace(0, 1, 100, dtype = x_train.dtype)
 153 y_test = phi(x_test)
 154
 155 for D in range(args.D_max + 1):
 156     fig = plt.figure()
 157
 158     ax = fig.add_subplot(1, 1, 1)
 159     ax.set_title(f'Degree {D}')
 160     ax.set_ylim(-0.1, 1.1)
 161     ax.plot(x_test, y_test, color = 'black', label = 'Test values')
 162     ax.scatter(x_train, y_train, color = 'blue', label = 'Train samples')
 163
 164     alpha = fit_alpha(x_train, y_train, D)
 165     ax.plot(x_test, pol_value(alpha, x_test), color = 'red', label = 'Fitted polynomial')
 166
 167     ax.legend(frameon = False)
 168
 169     fig.savefig(f'dd-example-{D:02d}.pdf', bbox_inches='tight')
 170
 171     plt.close(fig)
 172
 173 ######################################################################