Fit learning-rate and batch-size scaling laws from measured loss grids. For each model scale, select the lowest measured loss; ties favor the earlier learning-rate index, then the earlier batch-size index. Let Si be model scale at index i, ηi∗ its selected learning rate, and Bi∗ its selected batch size. Fit two independent, unweighted least-squares lines using natural logarithms:
logηi∗=logaη+bηlogSi. logBi∗=logaB+bBlogSi.Here, aη and aB are positive coefficients; bη and bB are exponents. Each residual is the observed log value minus its fitted log value. At target_scale, predict the coefficient multiplied by that scale raised to its exponent.
Return a Python dictionary with:
Input: model_scales = [1,4], learning_rates = [0.1,0.2], batch_sizes = [8,16], measured_losses = [[[1,2],[3,4]],[[4,3],[2,1]]], target_scale = 16
Output: {"optima":[{"model_scale":1,"learning_rate":0.1,"batch_size":8,"loss":1},{"model_scale":4,"learning_rate":0.2,"batch_size":16,"loss":1}],"learning_rate_coefficient":0.1,"learning_rate_exponent":0.5,"batch_size_coefficient":8,"batch_size_exponent":0.5,"learning_rate_residuals":[0,0],"batch_size_residuals":[0,0],"target_learning_rate":0.4,"target_batch_size":32}
Explanation: The two measured optima define exact square-root scaling for both hyperparameters.
Input: model_scales = [1,2], learning_rates = [0.1,0.2], batch_sizes = [8,16], measured_losses = [[[1,1],[1,1]],[[1,1],[1,1]]], target_scale = 4
Output: {"optima":[{"model_scale":1,"learning_rate":0.1,"batch_size":8,"loss":1},{"model_scale":2,"learning_rate":0.1,"batch_size":8,"loss":1}],"learning_rate_coefficient":0.1,"learning_rate_exponent":0,"batch_size_coefficient":8,"batch_size_exponent":0,"learning_rate_residuals":[0,0],"batch_size_residuals":[0,0],"target_learning_rate":0.1,"target_batch_size":8}
Input: model_scales = [1,2,4], learning_rates = [0.01,0.1,1], batch_sizes = [1,8,64], measured_losses.shape = (3, 3, 3), target_scale = 8
Output: {"optima":[{"model_scale":1,"learning_rate":0.01,"batch_size":1,"loss":1},{"model_scale":2,"learning_rate":0.1,"batch_size":64,"loss":1.01},{"model_scale":4,"learning_rate":1,"batch_size":8,"loss":1.02}],"learning_rate_coefficient":0.01,"learning_rate_exponent":3.321928,"batch_size_coefficient":2.828427,"batch_size_exponent":1.5,"learning_rate_residuals":[0,0,0],"batch_size_residuals":[-1.039721,2.079442,-1.039721],"target_learning_rate":10,"target_batch_size":64}
torch.argmin returns the first index when minima tie.
divmod recovers a row and column from a flattened index.
torch.linalg.lstsq solves a least-squares system.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Fit learning-rate and batch-size scaling laws from measured loss grids. For each model scale, select the lowest measured loss; ties favor the earlier learning-rate index, then the earlier batch-size index. Let Si be model scale at index i, ηi∗ its selected learning rate, and Bi∗ its selected batch size. Fit two independent, unweighted least-squares lines using natural logarithms:
logηi∗=logaη+bηlogSi. logBi∗=logaB+bBlogSi.Here, aη and aB are positive coefficients; bη and bB are exponents. Each residual is the observed log value minus its fitted log value. At target_scale, predict the coefficient multiplied by that scale raised to its exponent.
Return a Python dictionary with:
Input: model_scales = [1,4], learning_rates = [0.1,0.2], batch_sizes = [8,16], measured_losses = [[[1,2],[3,4]],[[4,3],[2,1]]], target_scale = 16
Output: {"optima":[{"model_scale":1,"learning_rate":0.1,"batch_size":8,"loss":1},{"model_scale":4,"learning_rate":0.2,"batch_size":16,"loss":1}],"learning_rate_coefficient":0.1,"learning_rate_exponent":0.5,"batch_size_coefficient":8,"batch_size_exponent":0.5,"learning_rate_residuals":[0,0],"batch_size_residuals":[0,0],"target_learning_rate":0.4,"target_batch_size":32}
Explanation: The two measured optima define exact square-root scaling for both hyperparameters.
Input: model_scales = [1,2], learning_rates = [0.1,0.2], batch_sizes = [8,16], measured_losses = [[[1,1],[1,1]],[[1,1],[1,1]]], target_scale = 4
Output: {"optima":[{"model_scale":1,"learning_rate":0.1,"batch_size":8,"loss":1},{"model_scale":2,"learning_rate":0.1,"batch_size":8,"loss":1}],"learning_rate_coefficient":0.1,"learning_rate_exponent":0,"batch_size_coefficient":8,"batch_size_exponent":0,"learning_rate_residuals":[0,0],"batch_size_residuals":[0,0],"target_learning_rate":0.1,"target_batch_size":8}
Input: model_scales = [1,2,4], learning_rates = [0.01,0.1,1], batch_sizes = [1,8,64], measured_losses.shape = (3, 3, 3), target_scale = 8
Output: {"optima":[{"model_scale":1,"learning_rate":0.01,"batch_size":1,"loss":1},{"model_scale":2,"learning_rate":0.1,"batch_size":64,"loss":1.01},{"model_scale":4,"learning_rate":1,"batch_size":8,"loss":1.02}],"learning_rate_coefficient":0.01,"learning_rate_exponent":3.321928,"batch_size_coefficient":2.828427,"batch_size_exponent":1.5,"learning_rate_residuals":[0,0,0],"batch_size_residuals":[-1.039721,2.079442,-1.039721],"target_learning_rate":10,"target_batch_size":64}
torch.argmin returns the first index when minima tie.
divmod recovers a row and column from a flattened index.
torch.linalg.lstsq solves a least-squares system.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
