Double Descent, Optimization, and Integration¶
Background¶
A while ago Advani and Saxe released High-dimensional dynamics of generalization error in neural networks That showed a theoretical model of the double desent phenomenon, where models begin to generalize better after the number of parameters in the model grows beyond the interpolation threshold. A few other groups went on to look into the phenomenon, and OpenAI made a pretty thorough post about the topic.
The most interesting thing though, is that this also applies to data when holding the number of parameters fixed. Essentially, the more data is added, the worse the validation performance gets. This is extremely counter intuitive, as we would imagine that so long as the train and test distributions are identical, more data should always be better.
One possible culprit might that we're treating the inference problem as an optimization problem rather than an integration problem, as would be suggested by bayesian inference.
To verify this, we'll train a models with optimizers and bayesian integrators on MNIST and compare their performance on the same test set as the amount of training data increases.
For the optimizers, we'll use plain SGD and ADAM. For the integrators, we'll use an implemenation of ATMC from Bayesian Inference for Large Scale Image Classification.
We'll look at the average performance of individual models and their ensembles as a function of the number of training datapoints for the MNIST dataset.
#Import dependencies
import math
import numpy as np
import torch as t
import torch.nn as nn
import torch.nn.functional as f
import torch.optim as optim
import matplotlib.pyplot as plt
from tqdm import tqdm
from torchvision.datasets import MNIST as MNIST
#Sampler taken from https://github.com/CurtisHuebner/ATMC
class ATMC(optim.Optimizer):
def __init__(self, params,h=0.001,d=1,m=1):
defaults = dict(h=h, m=m,d=d)
super(ATMC, self).__init__(params, defaults)
def step(self,closure=None):
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
h = group['h']
m = group['m']
d = group['d']
for param in group['params']:
if param.grad is None:
continue
g = param.grad
n = t.randn_like(g)
param_state = self.state[param]
if 'momentum' not in param_state:
param_state['momentum'] = t.zeros_like(g).detach()
p = param_state['momentum']
if 'regulator' not in param_state:
param_state['regulator'] = t.zeros_like(g).detach()
xi = param_state['regulator']
alpha = f.relu(d-xi)
beta = (alpha + xi)
#TODO:Replace with exact integration of the underlying DE
term = -g*h-beta*p*h+t.sqrt(2*h*alpha*m)*n
param.data += h*(p/m)
param_state['regulator'] += h*(p**2/m-1)
param_state['momentum'] += term
return loss
#Define a GLM classifier
FEATURES = t.randn(784,15)/math.sqrt(784)
class MNISTClassifier(nn.Module):
def __init__(self):
super(MNISTClassifier, self).__init__()
self.weights = nn.Linear(15,10)
self.features = FEATURES
self.register_buffer("random_features",self.features)
self.loss = nn.CrossEntropyLoss(reduction='none')
def forward(self,x,y):
x = x.view(-1,784)
h = t.matmul(x,self.features)
logits = self.weights(h)
loss = self.loss(logits,y)
return loss
def prior(self):
#return t.sum((self.weights.weight*math.sqrt(10))**2)/2 + t.sum(self.weights.bias**2)/2
return 0
def logits(self,x):
x = x.view(-1,784)
h = t.matmul(x,self.features)
logits = self.weights(h)
return logits
#Train the model
def train_model(model,opt,train_data,train_labels,test_data,test_labels):
for i in range(30000):
opt.zero_grad()
loss = t.sum(model(train_data,train_labels)) + model.prior()
loss.backward()
opt.step()
final_train_loss = t.mean(model(train_data,train_labels).detach())
final_test_loss = t.mean(model(test_data,test_labels).detach())
return final_train_loss,final_test_loss
#Get the dataset
def get_mnist():
PATH_TO_MNIST = "./"
mnist_ds = MNIST(PATH_TO_MNIST,download=True,train=True)
mnist = np.stack([np.array(x[0]) for x in mnist_ds])
labels = np.stack([np.array(x[1]) for x in mnist_ds])
mean = np.mean(mnist)
std = np.std(mnist)
scaled_mnist = (mnist-mean)/std
return scaled_mnist,labels
#Evaluate multiple models
def train_val_mnist(optimiser):
scaled_mnist,labels = get_mnist()
train_list = []
test_list = []
ens_list = []
n_list = []
for n in tqdm(range(1,300,10)):
N = 15
train_loss_mean = 0
test_loss_mean = 0
models = []
for i in range(N):
n_test = 1000
n_train = n
test_data = t.FloatTensor(scaled_mnist[:n_test])
test_labels = t.LongTensor(labels[:n_test])
train_data = t.FloatTensor(scaled_mnist[n_test:n_train+n_test])
train_labels = t.LongTensor(labels[n_test:n_train+n_test])
model = MNISTClassifier()
opt = optimiser(model.parameters())
final_train_loss,final_test_loss = train_model(model,opt,train_data,train_labels,test_data,test_labels)
train_loss_mean += final_train_loss/N
test_loss_mean += final_test_loss/N
models.append(model)
train_list.append(train_loss_mean)
test_list.append(test_loss_mean)
ens_list.append(ensemble(models,test_data,test_labels))
n_list.append(n)
return train_list,test_list,ens_list,n_list
def ensemble(models, test_data,test_labels):
log_p_vectors = []
for model in models:
logits = model.logits(test_data)
log_p_vectors.append(f.log_softmax(logits,dim=1).detach())
log_p = t.logsumexp(t.stack(log_p_vectors),dim=0) - math.log(len(models))
return f.nll_loss(log_p,test_labels)
opt = lambda x: ATMC(x, h=0.003)
bayes_train,bayes_test,bayes_ens,bayes_n = train_val_mnist(opt)
opt = lambda x: optim.SGD(x,lr=0.003)
sgd_train, sgd_test, sgd_ens, sgd_n = train_val_mnist(opt)
opt = lambda x: optim.Adam(x,lr=0.003)
adam_train, adam_test, adam_ens, adam_n = train_val_mnist(opt)
plt.title("Train")
plt.plot(bayes_n,bayes_train)
plt.plot(sgd_n,sgd_train)
plt.plot(adam_n,adam_train)
plt.show()
plt.title("Test")
plt.plot(bayes_n,bayes_test)
plt.plot(sgd_n,sgd_test)
plt.plot(adam_n,adam_test)
plt.show()
plt.title("Ensemble")
plt.plot(bayes_n,bayes_ens)
plt.plot(sgd_n,sgd_ens)
plt.plot(adam_n,adam_ens)
plt.show()
plt.title("ATMC")
plt.plot(bayes_n,bayes_ens)
plt.plot(bayes_n,bayes_test)
plt.show()
plt.title("SGD")
plt.plot(sgd_n,sgd_ens)
plt.plot(sgd_n,sgd_test)
plt.show()
plt.title("ADAM")
plt.plot(adam_n,adam_ens)
plt.plot(adam_n,adam_test)
Results¶
Overall, there are a few things that stand out:
-Adam is terrible before the interpolation threshold.
-Single samples from an MCMC sampler don't display the double desent phenomenon, but their ensemble does.
-Ensembling multiple SGD models together makes no difference.
-This seems to support the hypothesis that double desent is a property of the induced likelyhood function, rather than a product of the optimization + initialization procedure.
-For the MCMC sampler, It's important to note that the prior distribution being used is improper and so is the posterior past the interpolation threshold, thus without specifying a prior, bayesian inference is not really well defined.