Stabilize Double Poisson training procedure

[spencermyoung513]

In our brief experiments training neural networks with the Double Poisson NLL, we have observed a large variance in the resultant model's metrics (primarily MSE). This could be due to a variety of factors. The goal is to design a training process with the Double Poisson NLL that consistently outputs good models (with low variance in terms of model MSE across random initializations).

Areas to tinker with:

• [ ] Model initialization (different random initializations for model weights, along with schematic adjustments to the bias of the final layer to account for empirical statistical parameters)
• [ ] Model architecture (residual connections, batch normalization, etc.)
• [x] Optimization algorithm (Adam vs. AdamW vs. SGD, learning rate warmup / cosine annealing / other scheduling, etc.)
• [x] Training time (number of epochs)

[spencermyoung513]

Got some encouraging results with the following training settings:

epochs: 2000
batch_size: 32
lr: 1e-3
weight_decay: 1e-5
optimizer: AdamW
lr_scheduler: CosineAnnealingLR

I also trained each network 10 times. Crossing my fingers that we've found some stability.