Hyperparameter tuning

[forklady42]

Overview

Systematically tune hyperparameters to improve model performance beyond the current baseline.

Tasks

• Identify hyperparameters to tune (e.g., lr, n_residual_blocks, n_channels, warmup_length, batch size, optimizer settings)
• Choose a tuning strategy (e.g., grid search, random search, or Bayesian optimization via W&B sweeps)
• Run sweep and collect results
• Analyze results and select best configuration
• Update default config(s) with best-found hyperparameters and document findings

Relevant Config Parameters

See src/electrai/configs/ — key parameters include lr, epochs, nbatch, n_channels, n_residual_blocks, warmup_length.

Acceptance Criteria

• Tuning sweep completed and results logged
• Best hyperparameter configuration identified and documented
• Default config updated if improvement is significant

[forklady42]

Optuna set up in #84

[forklady42]

Because our input data's dimensions are only guaranteed to be divisible by 4 (2^2), we've been limited to a depth of two downsample layers. I experimented with a couple options for increasing depth to see if higher depths would be crucial for optimization. The initial options I tried were 1) to interpolate data to fit expected size when upsampling and 2) use dilated convolutions at the bottleneck to increase the receptive field and increase depth without downsampling. resunet-deeper below is option 1, and resunet-dilated is option 2.

After 15 epochs, the baseline is outperforming the deeper models on val_loss. Out of the experiments, the dilated approach consistently performs better than the interpolated model. The dilated model converges faster early on (lower train loss) but has a growing train-val gap — train is 0.0233 vs val 0.0383 at epoch 15. This suggests the extra parameters are starting to overfit.

I may try the dilated approach on a larger dataset (6000 structures) to see how it does at larger scale, but based on the initial performance, I'm not concerned about including depths >2 in the hyperparameter sweep.

[forklady42]

Depth & Dilated Bottleneck Experiment Summary (by Claude)

Tested whether increased model depth or dilated bottleneck convolutions improve over the baseline depth=2 ResUNet. Branch: betsy/dilated-bottleneck.

Approaches tested

1. Baseline — depth=2, 32 channels, no dilations
2. Dilated bottleneck (approach 3) — depth=2, 32 channels, bottleneck_dilations=[1, 2, 4] (3 dilated ResBlocks replacing the standard bottleneck, increasing receptive field without downsampling)
3. Deeper model (approach 4) — depth=3, 32 channels, no dilations (relies on trilinear interpolation for skip connection size mismatches when input isn't divisible by 8)

Dataset 2 results (variable input sizes, 16 epochs)

All runs: 32 channels, lr=0.01, gradient_clip=5.0, mp-resunet-ablation W&B project.

Epoch	Baseline (d=2) val	Dilated (d=2) val	Deeper (d=3) val
3	0.072	0.073	0.078
7	0.044	0.046	0.058
11	0.036	0.040	0.044
15	0.033	0.038	0.041

Dataset 4 results (128³ uniform inputs)

Attempted depth=3 on the uniform 128³ dataset. Had to reduce to 16 channels due to OOM at 32 channels on 4×A100 80GB, even with gradient checkpointing and mixed precision. Best run (resunet-128-dataset, lr=0.001):

Epoch	Train loss	Val loss
0	1.338	1.149
3	0.074	0.089
8	0.055	0.063
12	0.097	0.108

Loss diverged after epoch 9, likely due to numerical instability in the output normalization layer (division by near-zero output sums). Added an epsilon guard to address this for future runs.

Conclusions

• Baseline depth=2 wins — simplest model achieved the best val_loss (0.033) on dataset 2.
• Dilated bottleneck overfits — lowest train loss (0.023) but growing train-val gap (0.038 val). Extra parameters don't help generalization.
• Depth=3 converges slower and never catches up to baseline through 16 epochs.
• Depth=3 on 128³ data is memory-constrained (requires 16 channels instead of 32) and numerically unstable at higher learning rates.
• Greater depth does not appear to improve performance for this architecture. Recommend focusing tuning efforts on other hyperparameters.

https://wandb.ai/PrinceOA/mp-resunet-ablation

[forklady42]

Kernel Size Sweep Results (Tier 1)

Setup

• Study: resunet_hpo_tier1_kernel_v2 (30 trials, TPE sampler, median pruner)
• Search space: kernel_size ∈ {5, 7, 9}, depth ∈ {1, 2}, lr ∈ [0.001, 0.01] (log)
• Fixed: n_channels=32, n_residual_blocks=1 (best from prior sweeps)
• Hardware: 1× A100-SXM4-40GB, 15 epochs/trial
• W&B project: betsy-mp-hpo-tier1

Completed Trials (depth=2)

Trial	kernel_size	lr	val_loss
2	5	0.0083	0.0282
0	7	0.0093	0.0291
19	9	0.0073	0.0296
3	9	0.0076	0.0303

Completed Trials (depth=1)

Trial	kernel_size	lr	val_loss
1	7	0.0016	0.0343
4	9	0.0035	0.0350

Key Findings

1. kernel_size=5 remains the best (0.0282), confirmed by a second independent run (trial 20, pruned at epoch 12 with val=0.0314, on track for ~0.028).
2. kernel_size=7 is close (0.0291, ~3% worse) and kernel_size=9 is competitive (0.0296, ~5% worse). Larger kernels don't improve performance — the receptive field from depth=2 downsampling already captures sufficient spatial context.
3. depth=2 dominates depth=1 across all kernel sizes (~0.028–0.030 vs 0.034–0.035). No depth=1 configuration is competitive.
4. lr sweet spot is 0.007–0.009 (log scale). Trials with lr < 0.003 consistently underperformed or were pruned.

Pruner Analysis

17 of 24 trials were pruned by the median pruner. Analysis of intermediate values shows the pruner was overly aggressive — trial 20 (k=5, d=2, lr=0.0089) was pruned at epoch 12 with val=0.0314 while on an identical trajectory to the best trial (0.0282 at epoch 14). Future sweeps should consider increasing n_warmup_steps or using a less aggressive pruner.

Confirmed Tier 1 Hyperparameters

Parameter	Best Value	Confidence
n_channels	32	High
kernel_size	5	High (7 viable alternative)
n_residual_blocks	1	High
depth	2	High (with current data grid sizes; deeper requires padding fix)
lr	~0.007–0.009	High