Improve learning rate selection