Optimal Learning Rate Schedule for Balancing Effort and Performance

TL;DR

The paper presents a normative framework that casts learning-rate scheduling as an optimal-control problem balancing task performance against the learning cost. It yields a closed-loop optimal learning-rate rule that depends only on current and expected final performance, and it shows this rule generalizes across tasks and architectures. For simple models, an open-loop solution is derived, and the impact of discounting on learning speed is analyzed with analytical approximations and simulations. A biologically plausible episodic-memory mechanism is proposed to estimate future performance, enabling near-optimal control without full trajectory knowledge. The work links self-regulated learning, effort allocation, and memory-based performance estimation, with implications for both neuroscience and machine learning practice.

Abstract

Learning how to learn efficiently is a fundamental challenge for biological agents and a growing concern for artificial ones. To learn effectively, an agent must regulate its learning speed, balancing the benefits of rapid improvement against the costs of effort, instability, or resource use. We introduce a normative framework that formalizes this problem as an optimal control process in which the agent maximizes cumulative performance while incurring a cost of learning. From this objective, we derive a closed-form solution for the optimal learning rate, which has the form of a closed-loop controller that depends only on the agent's current and expected future performance. Under mild assumptions, this solution generalizes across tasks and architectures and reproduces numerically optimized schedules in simulations. In simple learning models, we can mathematically analyze how agent and task parameters shape learning-rate scheduling as an open-loop control solution. Because the optimal policy depends on expectations of future performance, the framework predicts how overconfidence or underconfidence influence engagement and persistence, linking the control of learning speed to theories of self-regulated learning. We further show how a simple episodic memory mechanism can approximate the required performance expectations by recalling similar past learning experiences, providing a biologically plausible route to near-optimal behaviour. Together, these results provide a normative and biologically plausible account of learning speed control, linking self-regulated learning, effort allocation, and episodic memory estimation within a unified and tractable mathematical framework.

Figures (5)

• Figure 1: Theoretical Framework: During neural network training, the optimal learning rate for weight updates is computed to maximize the cumulative internal reward throughout learning. The internal reward is the difference between the reward collected through performing a task $P(t)$ and the cost $C(\mu(t))$ of learning with the rate $\mu$.
• Figure 2: Closed-loop expression generalizes to arbitrary tasks and networks.(a,c): Optimal learning rate schedules for a two-layer non-linear network on a teacher-student setup trained with gradient flow (a), and digit classification/MNIST task trained with SGD (c). (b, d): Cumulative internal rewards for the optimal solution from \ref{['eq:closed_loop']} (grey), baseline agent with the best performing fixed learning rate (solid blue), and numerical optimization (purple dots). (e, f) Learning rates (e) and cumulative internal rewards (f) for different learning strategies. Darker tones represent strategies with more effort, colours represent strategy profiles (purple - flat rate, orange - learn-then-rest, blue - ramp-up, bold orange - optimal). (g) Plot of performance/cost trade-off for different learning strategies (each dot is a different learning rate schedule). The colourbar shows the cumulative internal reward rate, i.e. the optimization objective. The yellow star shows the optimal strategy. Strategies lie on the lines of learning profiles (colours as in (e,f)). (h) Cumulative rewards of the optimal learning strategy as a function of relative estimation error of final performance $\hat{P}(T)/P(T)$, spanning under- and over-confident regimes (colorbar shows total effort). Coloured dashed lines give the best-performing strategy for each learning profile, with the true optimum indicated by a star. Optimal profiles with estimation errors up to 20% exceed the best flat strategy (indicated by the vertical gray dashed lines). The same error bounds are marked by white lines on the effort colorbar. Experimental details are provided in section \ref{['sec:implementation_details']}.
• Figure 3: Estimating final performance through episodic memory.(a, b): Estimated future trajectory (orange) throughout learning, based on the observed trajectory $\mathcal{T}_{o}$ (blue), using similarities with trajectories stored in memory $\mathcal{T}_{h}$ (gray). True future trajectory shown in green. (c): Estimation error $\hat{P}(T|t)-P(T)$ as a function of within-episode time $t$. For each memory size (colour; number of past learning episodes used), we plot the 25th and 75th percentiles of the error across trials (smoothed). As memory increases and more of the current trajectory is observed, the error both contracts and becomes increasingly symmetric, indicating a reduction in bias as well as variance. (d): Learning-trajectory variability (measured as standard deviation) versus estimation error, coloured by the order of learning instance in memory. Higher-variance trajectories exhibit larger estimation errors. Later learning instances cluster near low variability and near-zero error, yet even at late stages increased trajectory variability remains associated with reduced predictability.
• Figure 4: Optimal learning strategies change with parameters in predictable ways.(a-c) A linear perceptron learning a regression task. (a) Optimal learning rate (dots: numerically optimized, line: analytical expression from Eq. \ref{['eq:lr_openloop_no_alpha']}). (b) Increasing the cost of control coefficient $\beta$ leads to lower learning rates in a non-linear way. (c) Task difficulty, defined as the norm of the difference between the target and initial weights $d = \| \mathbf w^* - \mathbf w_0\|^2$, increases optimal learning rate. (d,e) Learning rate for a one linear neuron learning a linear regression task with discounted time, an analytical approximation from Eq. \ref{['eq:gamma_approximation']} (e), or numerically evaluated (d). The approximation retains qualitative properties of the exact solution, showing that stronger discounting leads to lower learning rates. (f) Two-layer linear network learning a teacher-student regression task, showing a non-linear effect of changing $\gamma$ on the optimal learning rate. Overall, decreasing the cost of effort, increasing task difficulty or increasing the discount factor scales the optimal learning rate up. Furthermore, strong discounting can completely alter the shape of learning schedule when nonlinear learning dynamics is present. Experimental details are provided in section \ref{['sec:implementation_details']}.
• Figure 5: Comparison to data. Left: Instantaneous reward rate over time in a two-alternative forced choice task. Rats were grouped by reaction times relative to their individual baselines from prior experiments (fast, black; slow, purple). Data adapted from masis_strategically_2023, Fig. 7b. Right: Neural network model performing a two-Gaussian discrimination task. One model follows the optimal learning strategy (purple), while the other uses the same profile with a reduced learning-rate magnitude (black). Animals exhibit a trade-off between immediate reward rate and learning speed: lower immediate rewards, for example due to longer reaction times under greater control, can yield faster learning and higher cumulative reward. Our framework captures this trade-off as a balance between control cost and learning rate. Experimental details are provided in section \ref{['sec:implementation_details']}.