mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-09-09 00:54:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Multiply lr by 10; simplify Cain.

Browse Source
This commit is contained in:
Daniel Povey 2022-06-04 15:48:33 +08:00
parent 679972b905
commit a9a172aa69
2 changed files with 23 additions and 228 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

249
egs/librispeech/ASR/pruned_transducer_stateless7/optim.py
View File

@ -48,33 +48,13 @@ class Cain(Optimizer):
the model and "moved them into the optimizer". Now we allow
parameters to have any rms value (that's <= 10), but we:
(a) make the update speed for a parameter proportional to
its rms value, using target_rms (0.1) as the reference point
where we just use the standard Adam-type formula.
(b) simulate the effect of learning the weight_scale or bias_scale
in log space using standard Adam. (see scale_exp_avg_sq in the
its rms value.
(b) Multiply the learning rate by 10 because we are getting rid of
target_rms=0.1
(c) simulate the effect of learning the weight_scale or bias_scale
in log space using standard Adam, with a scale to prevent it
from learning too fast. (see scale_exp_avg_sq in the
code below).
(iv) Periodically (every 2000 batches) perform a "change of basis" on
all non-convolutional axes of each parameter tensor, and reset the
Adam statistics. This means an orthogonal change of co-ordinates
in the space corresponding to each dimension of a tensor, e.g. in
a 256 x 512 parameter tensor we'll have a 256x256 orthogonal matrix
and a 512x512 one. This does a better job of separating "important"
from "unimportant" directions in activation space, and empirically
improved convergence, final loss function and WERs.
(v) Change the Adam-type update, where the rms update value in each parameter
is the same, to a (diagonal) natural-gradient-type update, where if a
parameter has larger then average gradients, we give it a smaller
update. Because, in principle, this could lead to unbounded update
magnitudes, we limit these by introducing the "max_eff_lr" parameter,
which defines the maximum learning rate "if we had an Adam-type
update". So what we really implement is a cross between Adam and
natural gradient, where larger "max_eff_lr" means more natural-gradient-like
(i.e. more aggressive in directions with small gradients).
We recommend to set max_eff_lr to the initial learning rate ("lr")
in the learning rate schedule so that the update starts out
equivalent to Adam. Empirically this change to natural-gradient-type
update, epecially when combined with a faster-decaying learning rate
schedule, gives faster convergence.
The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
@ -84,14 +64,12 @@ class Cain(Optimizer):
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
ng_scale (float, optional): scale on natural-gradient-like update,
interpolating it with an Adam-type update.
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-8)
target_rms (float, optional): target root-mean-square value of
parameters, if they fall below this we will stop applying weight decay.
scale_speed (float, optional): we multiply the learning rate by this value
when learning (in log space) the scales of the parameter tensors
rms_eps (float, optional): epsilon value such that when parameter
rms value goes below this, we stop learning slower for that parameter,
i.e. we treat the rms as if it were rms_eps. In practice, rms values
@ -116,22 +94,16 @@ class Cain(Optimizer):
def __init__(
self,
params,
lr=1e-3,
ng_pow=-0.5,
ng_eps=0.1,
lr=1e-2,
betas=(0.9, 0.98),
eps=1e-8,
target_rms=0.1,
scale_speed=0.05,
rms_eps=1.0e-05,
rms_max=10.0,
):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= ng_eps < 1.0:
raise ValueError("Invalid natural gradient epsilon: {}".format(ng_eps))
if not -1.0 <= ng_pow < 1.0: # -1.0 would be actual natural gradient
raise ValueError("Invalid natural gradient power: {}".format(ng_pow))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))
if not 0.0 <= betas[0] < 1.0:
@ -142,8 +114,8 @@ class Cain(Optimizer):
raise ValueError(
"Invalid beta parameter at index 1: {}".format(betas[1])
)
if not 0 < target_rms <= 10.0:
raise ValueError("Invalid target_rms value: {}".format(target_rms))
if not 0 < scale_speed < 1.0:
raise ValueError("Invalid scale_speed value: {}".format(scale_speed))
if not 0.1 < rms_max <= 1000.0:
raise ValueError("Invalid rms_max value: {}".format(rms_max))
if not 0.0 < rms_eps < 0.1:
@ -151,11 +123,9 @@ class Cain(Optimizer):
defaults = dict(
lr=lr,
ng_eps=ng_eps,
ng_pow=ng_pow,
betas=betas,
eps=eps,
target_rms=target_rms,
scale_speed=0.05,
rms_eps=rms_eps,
rms_max=rms_max,
)
@ -180,9 +150,7 @@ class Cain(Optimizer):
for group in self.param_groups:
beta1, beta2 = group["betas"]
lr = group["lr"]
ng_eps = group["ng_eps"]
ng_pow = group["ng_pow"]
target_rms = group["target_rms"]
scale_speed = group["scale_speed"]
rms_eps = group["rms_eps"]
eps = group["eps"]
rms_max = group["rms_max"]
@ -221,7 +189,6 @@ class Cain(Optimizer):
dtype=p.dtype)
delta, exp_avg_sq = state["delta"], state["exp_avg_sq"]
step = state["step"]
@ -250,7 +217,7 @@ class Cain(Optimizer):
scale_denom = (scale_exp_avg_sq.sqrt()).add_(eps)
scale_alpha = -lr * (scale_bias_correction2 ** 0.5) * (1-beta1)
scale_alpha = -lr * (scale_bias_correction2 ** 0.5) * scale_speed * (1-beta1)
# scale_delta is going to be the change in log-scale (before momentum), i.e. if
# p = underlying_param * scale.exp(),
@ -270,46 +237,19 @@ class Cain(Optimizer):
scale_delta)
delta.add_(p, alpha=scale_delta)
grad = self._change_coordinates(grad, state, forward=True)
p_transformed = self._change_coordinates(p, state, forward=True, is_grad=False)
# Decay the second moment running average coefficient
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
bias_correction2 = 1 - beta2 ** self._step_for_bias_correction(step)
bias_correction2 = 1 - beta2 ** step
grad_rms = (exp_avg_sq.sqrt()).add_(eps)
grad_rms = (exp_avg_sq.sqrt()).add_(eps).mul_(1.0 / bias_correction2)
this_delta = grad / grad_rms
if p.ndim > 1:
grad_rms_mean = grad_rms.mean()
smoothed_grad_rms = ng_eps * grad_rms_mean + grad_rms
lr_factor = (smoothed_grad_rms ** ng_pow) + ng_eps * (grad_rms_mean ** ng_pow)
lr_factor_mean = lr_factor.mean()
# this `lr` contains all the terms in learning rate except param_rms. Imagine
# that param_rms is 1.0 for purposes of computing the shrinkage coefficient; actually
# it does not matter.
this_lr = (lr / (target_rms * lr_factor_mean)) * lr_factor
else:
# turn off this natural gradient for biases, it seems to have
# a bad effect, possibly because we don't decorrelate in this case, or
# because their gradients are not dominated by noise.
this_lr = lr / target_rms
this_lr = state["param_rms"].clamp(min=rms_eps) * lr * bias_correction2
# change sign so it's an actual parameter change,
# incorporate param_rms at this point.
this_delta *= -this_lr * state["param_rms"].clamp(min=rms_eps)
# The shrinkage coefficint necessary to bring data with variance=1.0 + lr**2
# back down to data with variance=1.0, is: (1 + lr**2) ** -0.5, which means
# multiplying by about (1 - 0.5 * lr**2).
# we then take the -0.5 * lr**2 * (existing_param) as the "delta" (change in param)
# that comes from this update.
shrinkage_delta = p_transformed * (-0.5 * (this_lr**2))
this_delta += shrinkage_delta
this_delta = self._change_coordinates(this_delta, state, forward=False)
alpha = (1-beta1)
alpha = -(1-beta1) * this_lr
# below, will do: delta.add_(this_delta, alpha=alpha)
else:
@ -329,150 +269,6 @@ class Cain(Optimizer):
return loss
def _change_coordinates(self,
grad: Tensor,
state: dict,
forward: bool,
is_grad: bool = True):
"""
Possibly performs some magnitude-preserving changes of co-ordinates on `grad`.
If is_grad == False we skip any possible update or stats-accumulation steps.
"""
ndim = grad.ndim
numel = grad.numel()
for dim in range(grad.ndim):
# see for each dim in turn whether we want to perform any changes in co-ordinates,
# or store any stats.
size = grad.shape[dim]
if size <= 3 or (size % 2 == 1 and size < 128) or size >= 2048 or size == numel:
# we don't do any such co-ordinate changes in dims with sizes
# that are too small (no point) or large (too slow), or that are
# assumed convolutional (because they are odd and not too huge).
# We can revisit this later.
continue
grad = self._change_coordinates_for_dim(grad, state, dim, forward, is_grad)
return grad
def _change_coordinates_for_dim(self,
grad: Tensor,
state: dict,
dim: int,
forward: bool,
is_grad: bool,
orig_dim: int = -1):
assert grad.ndim > 1
if not (grad.ndim == 2 and dim == 1):
# normalize the format and recurse.
dims_order = []
rev_dims_order = []
ndim = grad.ndim
for i in range(dim):
dims_order.append(i)
rev_dims_order.append(i)
rev_dims_order.append(ndim-1)
for i in range(dim+1, ndim):
dims_order.append(i)
rev_dims_order.append(i-1)
dims_order.append(dim)
# e.g. ndim=4, dim=1, dims_order=(0,2,3,1), rev_dims_order=(0,3,1,2)
new_grad = grad.permute(*dims_order)
assert new_grad.shape[-1] == grad.shape[dim]
new_shape = new_grad.shape
new_grad = new_grad.reshape(-1, new_grad.shape[-1])
new_grad = self._change_coordinates_for_dim(new_grad, state, 1,
forward, is_grad,
orig_dim=dim)
return new_grad.reshape(new_shape).permute(*rev_dims_order)
# OK: grad.ndim == 2 and dim == 1
size = grad.shape[1]
if orig_dim == -1:
orig_dim = dim
step = state["step"]
must_store_stats, must_zero_stats = self._must_store_stats(step)
if must_store_stats and forward and is_grad:
# store stats for 200 iters preceding when we estimate the
# transform.
stats_name = f"cov_{orig_dim}"
if not stats_name in state:
state[stats_name] = torch.zeros(size, size, dtype=grad.dtype,
device=grad.device)
cov = state[stats_name]
if must_zero_stats:
cov.zero_()
cov += torch.matmul(grad.t(), grad) * (1/grad.shape[0])
transform_name = f"t_{orig_dim}"
if transform_name not in state:
# make sure they're all allocated at the start, may be better for
# fragmention and debugging reasons (although may not).
state[transform_name] = torch.eye(size, device=grad.device,
dtype=grad.dtype)
must_estimate_transform = self._must_estimate_transform(step)
if must_estimate_transform and forward and is_grad:
stats_name = f"cov_{orig_dim}"
cov = state[stats_name]
l, U = cov.symeig(eigenvectors=True)
#print("estimate, l = ", l[::20])
t = U.t() # t is indexed (new_dim, old_dim) a.k.a. (transformed_dim, real_dim)
state[transform_name][:] = t
state["exp_avg_sq"].zero_() # zero exp-avg-sq stats, we'll restart these from scratch.
t = state[transform_name]
if forward:
return torch.matmul(grad, t.t())
else:
return torch.matmul(grad, t)
def _must_store_stats(self, step: int) -> Tuple[bool, bool]:
"""
Returns (must_store_stats, must_zero_stats), where
must_store_stats: bool, is True for the 100 or 200 steps preceding
a step on which we will estimate the transform.
must_zero_stats: bool, is True for only the 1st of the 100 or 200
steps mentioned above.
"""
if step < 4000:
if step < 2000:
return (step % 500 >= 400, step % 500 == 400)
else:
return (step % 1000 >= 800, step % 1000 == 800)
else:
return (step % 2000 >= 1800, step % 2000 == 1800)
def _must_estimate_transform(self, step: int) -> bool:
"""
Returns True if we will re-estimate the transform on this step
"""
if step < 4000:
if step < 2000:
return step % 500 == 0 and step > 0
else:
return step % 1000 == 0 and step > 0
else:
return step % 2000 == 0
def _step_for_bias_correction(self, step: int) -> bool:
"""
Return a modified step for bias-correction (actually just to give us the
right denominator for the exp-avg-sq stats).. this needs to take into
account how long it has been since we re-estimated the transforms, because
we zero the exp-avg-sq stats at those times.
The + 1 is for the "current step". We post-increment the step so it
starts from 0.
"""
if step < 4000:
if step < 2000:
return step % 500 + 1
else:
return step % 1000 + 1
else:
return step % 2000 + 1
class LRScheduler(object):
"""
@ -797,7 +593,7 @@ def _test_eve_cain():
input_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp()
output_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp()
for iter in [2,1,0]:
for iter in [1,0]:
fix_random_seed(42)
Linear = torch.nn.Linear if iter == 0 else ScaledLinear
m = torch.nn.Sequential(Linear(E, 200),
@ -808,8 +604,7 @@ def _test_eve_cain():
torch.randn(B, T, E, device=device, dtype=dtype) * output_magnitudes) for _ in range(20) ]
if iter == 0: optim = Eve(m.parameters(), lr=0.003)
elif iter == 1: optim = Cain(m.parameters(), lr=0.003)
elif iter == 2: optim = Cain(m.parameters(), lr=0.003, ng_pow=-0.5)
elif iter == 1: optim = Cain(m.parameters(), lr=0.03)
scheduler = Eden(optim, lr_batches=200, lr_epochs=10, verbose=False)
start = timeit.default_timer()

2
egs/librispeech/ASR/pruned_transducer_stateless7/train.py
View File

@ -207,7 +207,7 @@ def get_parser():
parser.add_argument(
"--initial-lr",
type=float,
default=0.003,
default=0.03,
help="The initial learning rate. This value should not need "
"to be changed.",
)