# Why does my training explode?

For simplicity purposes, I've changed the training data. The network should predict upper and lower bound on when an exponential function would cross a threshold, given a certain number of noisy samples in known intervals.

Here's my implementation:

First, here's the toy data generation method:

def gen_data(n: int) -> (torch.Tensor, torch.Tensor):
t = torch.linspace(0, 20, NUM_FEATURES).repeat(n, 1)
a = torch.rand((n, 1))
lim = torch.max(
torch.hstack([
torch.normal(0.2, 0.1, (n, 1)),
0.01 * torch.ones(n, 1)
]),
-1
).values

x = torch.exp(-a * t) + torch.normal(0, 0.01, t.size())
y = -torch.log(lim) / a.T
return x, y[0, :]


Now here's the loss implementation as per the linked article:

def huber(x):
torch.abs(x) < D,
(x ** 2) / 2,
D * torch.abs(x) - (D ** 2) / 2
)

def lube_loss(preds, truth):
lower, upper = preds.T
l1 = torch.abs(truth - (lower + upper) / 2)
l2 = torch.abs(upper - lower)
l3 = torch.where(
torch.logical_and(lower < truth, truth < upper),
torch.zeros_like(truth),
l1 + l2/2
)
return torch.mean(huber(K1 * l1) + huber(K2 * l2 + K3 * l3))


And here's the train code:

def init_weights(m: torch.nn.Module):
if not hasattr(m, 'weight'):
return
n = len(m.weight)
torch.nn.init.normal_(m.weight, 0, 1 / math.sqrt(n))

def nn_unit(idx: int, inputs: int, outputs: int) -> torch.nn.Module:
(f'h{idx}', torch.nn.Linear(inputs, outputs)),
(f'a{idx}', torch.nn.LeakyReLU()),
(f'drop{idx}', torch.nn.Dropout())]))

def train():
dataset = torch.utils.data.TensorDataset(*gen_data(100000))
dataset, batch_size=BATCH_SIZE, shuffle=True)
net = torch.nn.Sequential(OrderedDict([
('u1', nn_unit(1, NUM_FEATURES, NUM_FEATURES * 4)),
('u2', nn_unit(2, NUM_FEATURES * 4, NUM_FEATURES * 8)),
('u3', nn_unit(3, NUM_FEATURES * 8, NUM_FEATURES * 4)),
('u4', nn_unit(4, NUM_FEATURES * 4, NUM_FEATURES * 4)),
('u5', nn_unit(5, NUM_FEATURES * 4, NUM_FEATURES * 2)),
('u6', nn_unit(6, NUM_FEATURES * 2, NUM_FEATURES * 2)),
('out', torch.nn.Linear(NUM_FEATURES * 2, 2)),
]))
net.apply(init_weights)
optimizer = torch.optim.SGD(
net.parameters(), lr=1e-9, momentum=0.5, weight_decay=10
)
running_loss = 0
for i in range(1000000):
preds = net(x)
loss = lube_loss(preds, y)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
print('[%5d] loss: %.3f' %
(i + 1, running_loss / BATCH_SIZE))
running_loss = 0

lower, upper = net(X[:10, :]).T
for l, u, y in zip(lower, upper, Y):
print(f'{y}: [{l} - {u}]')

return net


And for completeness the imports and constants:

import math
import torch
import torch.utils.data
from collections import OrderedDict

NUM_FEATURES = 10
K1 = 2
K2 = 1
K3 = 4
D = 0.5
BATCH_SIZE = 100


The train function initially lowers the loss gradually, but then it starts to explode... I've tried lowering the learning rate (started at 0.1 and got to 1e-9), I've tried adding weight decay (started without it, and got up to 10). I've started with a shallow net with just one hidden layer and grew it (hoping it would be able to fit better with more layers and larger layers), it always ended up exploding... What am I missing?