PyTorch learning rate comparison in a marimo notebook

An interactive marimo notebook compares learning rate schedules for a PyTorch neural network on MNIST 1D. The tutorial starts with a linear regression loss landscape: small steps make slow progress, while oversized steps jump across the minimum. Sliders let the presenter change the learning rate and inspect the resulting path. The notebook cells react to those changes, updating the charts while the same Python notebook is displayed as an app.

The experiment compares a constant rate, step decay, and schedules with one, two or three cycles. In the initial 40-epoch runs, cyclical schedules lag during their gradual ramp-up. At a maximum rate of 0.05, the presenter reports faster convergence across all schedules. Raising it to 0.5 causes performance to collapse; slower ramp-up delays the failure but does not prevent it. Extending training to 100 epochs also changes the ramp-up period, which complicates direct comparison.

The presenter attributes the collapse partly to ReLU's zero-gradient region, while acknowledging that the notebook does not visually prove that explanation. Switching to sigmoid allows the cyclical methods to recover and converge in the demonstrated run. Constant and step-decay schedules still struggle with the oversized steps. The lesson is specific to this experiment: cycling can provide opportunities to recover from a poor learning rate choice, but the schedule and activation function both affect the outcome.