mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-09-19 05:54:20 +00:00
2324 lines
103 KiB
Python
2324 lines
103 KiB
Python
# Copyright 2022 Xiaomi Corp. (authors: Daniel Povey)
|
|
#
|
|
# See ../LICENSE for clarification regarding multiple authors
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from collections import defaultdict
|
|
from typing import List, Optional, Union, Tuple, List
|
|
from lhotse.utils import fix_random_seed
|
|
import torch
|
|
from scaling import ActivationBalancer
|
|
import random
|
|
from torch import Tensor
|
|
from torch.optim import Optimizer
|
|
import logging
|
|
|
|
|
|
class BatchedOptimizer(Optimizer):
|
|
"""
|
|
This class adds to class Optimizer the capability to optimize parameters in batches:
|
|
it will stack the parameters and their grads for you so the optimizer can work
|
|
on tensors with an extra leading dimension.
|
|
|
|
Args:
|
|
params:
|
|
"""
|
|
def __init__(self, params, defaults):
|
|
super(BatchedOptimizer, self).__init__(params, defaults)
|
|
|
|
|
|
def batched_params(self, param_group):
|
|
"""
|
|
This function returns an iterator over tuples (p, state), where
|
|
p is a `fake` parameter that is stacked (over axis 0) from real parameters
|
|
that share the same shape, and its gradient is also stacked;
|
|
`state` is the state corresponding to this fake parameter
|
|
(it will be physically located in the "state" for one of the real
|
|
parameters, the last one that has any particular shape and dtype).
|
|
|
|
The idea is, instead of doing:
|
|
<code>
|
|
for p in group["params"]:
|
|
state = self.state[p]
|
|
...
|
|
</code>
|
|
you can do:
|
|
<code>
|
|
for p, state in self.batched_params(group["params"]):
|
|
...
|
|
</code>
|
|
|
|
Args:
|
|
group: a parameter group, which is a list of parameters; should be
|
|
one of self.groups.
|
|
"""
|
|
batches = defaultdict(list) # `batches` maps from tuple (dtype_as_str,*shape) to list of nn.Parameter
|
|
|
|
for p in param_group:
|
|
assert p.grad is not None and "All parameters must have a grad when you call optimizer.step()"
|
|
assert not p.grad.is_sparse and "Sparse gradients not supported."
|
|
key = (str(p.dtype), *p.shape)
|
|
batches[key].append(p)
|
|
|
|
for p in param_group:
|
|
key = (str(p.dtype), *p.shape)
|
|
batch = batches[key]
|
|
if p is batch[0]: # if this is the 1st param in the batch...
|
|
# we arbitrarily store the state in the
|
|
# state corresponding to the 1st parameter in the
|
|
# group. class Optimizer will take care of saving/loading state.
|
|
state = self.state[p]
|
|
p_stacked = torch.stack(batch)
|
|
grad = torch.stack([p.grad for p in batch])
|
|
p_stacked.grad = grad
|
|
yield p_stacked, state # <-- calling code will do the actual optimization here!
|
|
# Now un-stack the parameter changes
|
|
for i,p in enumerate(batch):
|
|
p.copy_(p_stacked[i])
|
|
|
|
|
|
|
|
|
|
class PrAdam(BatchedOptimizer):
|
|
"""
|
|
Implements 'Projected Adam', a variant of Adam with projections for each
|
|
dim of each tensor.
|
|
|
|
|
|
Args:
|
|
params: The parameters or param_groups to optimize (like other Optimizer subclasses)
|
|
lr: The learning rate. We will typically use a learning rate schedule that starts
|
|
at 0.03 and decreases over time, i.e. much higher than other common
|
|
optimizers.
|
|
betas: beta1,beta2 are momentum constants for regular momentum, and moving sum-sq grad only for
|
|
scalar tensors. Must satisfy 0 < beta <= beta2 < 1.
|
|
size_lr_scale: A scaling factor on the learning rate, that we use to update the
|
|
scale of each parameter tensor. If each parameter were decomposed
|
|
as p * p_scale.exp(), where (p**2).mean().sqrt() == 1.0, size_lr_scale
|
|
is the scaling factor on the learning rate of p_scale.
|
|
param_rms_smooth0: Limiting value of smoothing proportion for parameter matrix, as
|
|
assumed rank of param covariance [==product of sizes on the other
|
|
tensor dims] approaches 0.
|
|
param_rms_smooth1: Smoothing proportion for parameter matrix, if assumed rank of
|
|
param covariance equals the dimension of the covaraince matrix.
|
|
param_rms_smooth{0,1} determine the smoothing proportions for other
|
|
conditions.
|
|
cov_min,cov_max: [IMPORTANT] 4-tuples of minimums and maximums of the diagonal values of
|
|
covariance matrices, after normalizing to unit-mean. The first 3 are
|
|
for smoothing the parameter covariance, normalized in 3 different ways:
|
|
(1) relative to its own diagonal (in a basis that diagonalizes the grad
|
|
covariance);
|
|
(2) multiplied by the grad covariance,
|
|
(3) in the canonical basis.
|
|
|
|
(4) is for smoothing the grad covariance used for (2)
|
|
cov_pow: This was mainly added for development and experimentation purposes;
|
|
it allows you to smooth the parameter covariance matrices at the
|
|
stages (1), (2), (3) of smoothing mentioned above, and also
|
|
the gradient covariance matrix used in stage (2) of smoothing. If the
|
|
1st 3 values are not 1.0 it will cause a measurable speed penalty because it
|
|
requires SVD. Recommend to leave all these at 1.0.
|
|
eps: A general-purpose epsilon to prevent division by zero
|
|
param_min_rms: Minimum root-mean-square value of parameter tensor, for purposes of
|
|
learning the scale on the parameters (we'll constrain the rms of each non-scalar
|
|
parameter tensor to be >= this value)
|
|
param_max_rms: Maximum root-mean-square value of parameter tensor, for purposes of
|
|
learning the scale on the parameters (we'll constrain the rms of each non-scalar
|
|
parameter tensor to be <= this value)
|
|
scalar_max: Maximum absolute value for scalar parameters (applicable if your
|
|
model has any parameters with numel() == 1)
|
|
size_update_period: The periodicity, in steps, with which we update the size (scale)
|
|
of the parameter tensor. This is provided to save a little time
|
|
in the update.
|
|
lr_update_period: [IMPORTANT]: A 2-tuple of ints that Determines the periodicity, in steps, with
|
|
which we update the learning-rate matrices. The first number is the periodicity at
|
|
the start of training, the second number is the periodicity
|
|
later in training, and we gradually increase from one to the other.
|
|
The reason for such a complicated schedule is that updating the learning
|
|
rate matrices is very slow, principally because it requires SVD, and SVD
|
|
seems to have quite slow implementations.
|
|
max_block_size: [IMPORTANT] The maximum block size in block-diagonal co-ordinate
|
|
transformations. You can probably set this to 512 or 1024. Larger
|
|
values will require more MEMORY and may be a bit slower, but should
|
|
lead to better optimization performance.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
params,
|
|
lr=3e-02,
|
|
betas=(0.9, 0.98),
|
|
size_lr_scale=0.1,
|
|
cov_min=(0.025, 0.0025, 0.02, 0.0001),
|
|
cov_max=(10.0, 80.0, 5.0, 400.0),
|
|
cov_pow=(1.0, 1.0, 1.0, 1.0),
|
|
param_rms_smooth0=0.4,
|
|
param_rms_smooth1=0.2,
|
|
eps=1.0e-08,
|
|
param_min_rms=1.0e-05,
|
|
param_max_rms=2.0,
|
|
scalar_max=2.0,
|
|
size_update_period=4,
|
|
lr_update_period=(200, 1000),
|
|
# grad_cov_period does not affect speed very much. If __name__ ==
|
|
# "__main__", keeping this coprime to 20 which is the number of
|
|
# sample pairs used in the test code
|
|
grad_cov_period=(3 if __name__ == "__main__" else 5),
|
|
max_block_size=1024,
|
|
):
|
|
|
|
|
|
|
|
defaults = dict(
|
|
lr=lr,
|
|
size_lr_scale=size_lr_scale,
|
|
cov_min=cov_min,
|
|
cov_max=cov_max,
|
|
cov_pow=cov_pow,
|
|
param_rms_smooth0=param_rms_smooth0,
|
|
param_rms_smooth1=param_rms_smooth1,
|
|
betas=betas,
|
|
eps=eps,
|
|
param_min_rms=param_min_rms,
|
|
param_max_rms=param_max_rms,
|
|
scalar_max=scalar_max,
|
|
size_update_period=size_update_period,
|
|
lr_update_period=lr_update_period,
|
|
grad_cov_period=grad_cov_period,
|
|
max_block_size=max_block_size,
|
|
)
|
|
|
|
super(PrAdam, self).__init__(params, defaults)
|
|
|
|
def __setstate__(self, state):
|
|
super(PrAdam, self).__setstate__(state)
|
|
|
|
|
|
@torch.no_grad()
|
|
def step(self, closure=None):
|
|
"""Performs a single optimization step.
|
|
|
|
Arguments:
|
|
closure (callable, optional): A closure that reevaluates the model
|
|
and returns the loss.
|
|
"""
|
|
loss = None
|
|
if closure is not None:
|
|
with torch.enable_grad():
|
|
loss = closure()
|
|
|
|
batch = True
|
|
for group in self.param_groups:
|
|
for p, state in self.batched_params(group["params"]):
|
|
|
|
# Perform optimization step
|
|
grad = p.grad
|
|
if grad.is_sparse:
|
|
raise RuntimeError(
|
|
"PrAdam optimizer does not support sparse gradients"
|
|
)
|
|
# State initialization
|
|
if len(state) == 0:
|
|
self._init_state(group, p, state)
|
|
self._step_one_batch(group, p, state)
|
|
|
|
|
|
return loss
|
|
|
|
|
|
def _init_state(self,
|
|
group: dict,
|
|
p: Tensor,
|
|
state: dict):
|
|
"""
|
|
Initializes state dict for parameter 'p'. Assumes that dim 0 of tensor p
|
|
is actually the batch dimension, corresponding to batched-together
|
|
parameters of a given shape.
|
|
|
|
Args:
|
|
group: Dict to look up configuration values.
|
|
p: The parameter that we are initializing the state for, actually
|
|
will be stacked-together "real" parameters
|
|
state: Dict from string to whatever state we are initializing
|
|
"""
|
|
eps = group["eps"]
|
|
size_update_period = group["size_update_period"]
|
|
max_block_size = group["max_block_size"]
|
|
|
|
state["step"] = 0
|
|
|
|
kwargs = {'device':p.device, 'dtype':p.dtype}
|
|
|
|
# 'delta' implements conventional momentum. There are
|
|
# several different kinds of update going on, so rather than
|
|
# compute "exp_avg" like in Adam, we store and decay a
|
|
# parameter-change "delta", which combines all forms of
|
|
# update. this is equivalent to how it's done in Adam,
|
|
# except for the first few steps.
|
|
state["delta"] = torch.zeros_like(
|
|
p, memory_format=torch.preserve_format
|
|
)
|
|
|
|
batch_size = p.shape[0]
|
|
numel = p.numel() // batch_size
|
|
|
|
if numel > 1:
|
|
# "param_rms" just periodically records the scalar root-mean-square value of
|
|
# the parameter tensor.
|
|
# it has a shape like (batch_size, 1, 1, 1, 1)
|
|
param_rms = _mean(p**2, exclude_dims=[0], keepdim=True).sqrt() + eps
|
|
state["param_rms"] = param_rms
|
|
|
|
state["scale_exp_avg_sq"] = torch.zeros_like(param_rms)
|
|
state["scale_grads"] = torch.zeros(size_update_period, *param_rms.shape,
|
|
**kwargs)
|
|
|
|
# exp_avg_sq is in the projected space. It is periodically zeroed, and we
|
|
# set zero_step whenever this happens.
|
|
state["exp_avg_sq"] = torch.zeros_like(
|
|
p, memory_format=torch.preserve_format
|
|
)
|
|
|
|
# ignore_rank1_dims = True means we don't do our basis-changing update
|
|
# on tensors like biases that only have one non-trivial dimension.
|
|
# "rank1" refers to the rank of our estimate of the parameter
|
|
# covariance, if we were to estimate it just from the current parameter
|
|
# value.
|
|
ignore_rank1_dims = True
|
|
|
|
trivial_update = True
|
|
for dim in range(1, p.ndim):
|
|
size = p.shape[dim]
|
|
if size == 1 or (size == numel and ignore_rank1_dims):
|
|
continue
|
|
trivial_update = False
|
|
|
|
state["trivial_update"] = trivial_update # so we know whether to zero exp_avg_sq stats.
|
|
if trivial_update:
|
|
return
|
|
|
|
# "zero_step" being a member of state is the sign that this parameter has
|
|
# at least one dim that has a projection. It also records the most recent
|
|
# step on which we zeroed state["exp_avg_sq"]
|
|
state["zero_step"] = 0
|
|
# last_param_scale_update records the last time we updated the part of the learning rate
|
|
# matrices that relates to the parameter covariance; we avoid doing this too often
|
|
# as it doesn't change very rapidly and the SVD is quite slow.
|
|
state["last_param_scale_update"] = -1
|
|
|
|
for dim in range(1, p.ndim):
|
|
size = p.shape[dim]
|
|
if size == 1 or (size == numel and ignore_rank1_dims):
|
|
continue
|
|
|
|
# Most of the time, (num_blocks, block_size) will equal (1, size)
|
|
num_blocks, block_size = self._get_block_size(size, max_block_size)
|
|
|
|
# Q_{dim} is be the learning-rate matrix for this
|
|
# dimension, a matrix of indexed [batch_idx, block_idx, diagonalized_coordinate, canonical_coordinate].
|
|
# this is used twice in the update step, once transposed and once without transpose.
|
|
Q = torch.eye(block_size, block_size, **kwargs).unsqueeze(0).unsqueeze(0).expand(
|
|
batch_size, num_blocks, block_size, block_size).contiguous()
|
|
state[f"Q_{dim}"] = Q
|
|
# cov_{dim} is the averaged-over-time covariance of parameters on this dimension, treating
|
|
# all other dims as a batch axis. Also initialize as identity.
|
|
state[f"param_cov_{dim}"] = Q.clone()
|
|
|
|
# grad_cov_{dim} is the covariance of gradients on this axis (without
|
|
# any co-ordinate changes), treating all other axes as as a batch axis.
|
|
# This is needed when we re-diagonalize, and to compute the
|
|
# grad_rms_{dim}. We store it as a decaying average, decaying with beta2
|
|
|
|
# only for purposes of computing the scalar factor on the
|
|
# learning rate; and, as a result of this, also contributes something
|
|
# to the gradient w.r.t. f"Q_{dim}", which is one reason
|
|
# why we allocate a variable to keep track of its moving average
|
|
# instead of just using a temporary and smoothing the scalar factor.
|
|
state[f"grad_cov_{dim}"] = torch.zeros_like(Q)
|
|
|
|
|
|
|
|
def _get_block_size(self, size: int, max_block_size: int) -> Tuple[int,int]:
|
|
"""
|
|
Returns information about the block size for a block-diagonal structure
|
|
of covariances and co-ordinate transformations.
|
|
|
|
size: The size of the parameter tensor on one of its dimensions, e.g.
|
|
a channel dim or kernel size
|
|
max_block_size: The maximum block size allowed (of the blocks in a
|
|
block-diagonal structure)
|
|
Returns:
|
|
(num_blocks, block_size)
|
|
where nun_blocks * block_size == size and block_size <= max_block_size
|
|
"""
|
|
for num_blocks in range(1, size):
|
|
block_size = size // num_blocks
|
|
if block_size * num_blocks == size and block_size <= max_block_size:
|
|
return (num_blocks, block_size)
|
|
assert False, (size, max_block_size) # code error or e.g. negative or
|
|
# zero inputs
|
|
|
|
|
|
def _step_one_batch(self,
|
|
group: dict,
|
|
p: Tensor,
|
|
state: dict):
|
|
"""
|
|
Do the step for one parameter, which is actually going to be a batch of
|
|
`real` parameters, with dim 0 as the batch dim.
|
|
Args:
|
|
group: dict to look up configuration values
|
|
p: parameter to update (actually multiple parameters stacked together
|
|
as a batch)
|
|
state: state-dict for p, to look up the optimizer state
|
|
"""
|
|
lr = group["lr"]
|
|
size_update_period = group["size_update_period"]
|
|
grad_cov_period = group["grad_cov_period"]
|
|
eps = group["eps"]
|
|
beta1 = group["betas"][0]
|
|
|
|
grad = p.grad
|
|
step = state["step"]
|
|
delta = state["delta"]
|
|
|
|
delta.mul_(beta1)
|
|
batch_size = p.shape[0]
|
|
numel = p.numel() // batch_size
|
|
if numel > 1:
|
|
# Update the size/scale of p, and set param_rms
|
|
scale_grads = state["scale_grads"]
|
|
scale_grads[step % size_update_period] = _sum(p * grad,
|
|
exclude_dims=[0],
|
|
keepdim=True)
|
|
if step % size_update_period == size_update_period - 1:
|
|
param_rms = state["param_rms"] # shape: (batch_size, 1, 1, ..)
|
|
param_rms.copy_(_mean(p ** 2, exclude_dims=[0],
|
|
keepdim=True).sqrt().clamp_(min=eps))
|
|
if step > 0:
|
|
# self._size_update() learns the overall scale on the
|
|
# parameter, by shrinking or expanding it.
|
|
self._size_update(group, scale_grads, p, state)
|
|
|
|
|
|
if numel == 1:
|
|
# For parameters with 1 element we just use regular Adam.
|
|
# Updates delta.
|
|
self._step_scalar(group, p, state)
|
|
else:
|
|
if self._is_lr_update_step(group, state):
|
|
self._update_param_cov(group, p, state)
|
|
|
|
P_proj = self._compute_bases(group, p.shape, state)
|
|
# Only update the parameter-dependent part of the learning
|
|
# rate matrices at most every other time we reach here, and
|
|
# less frequently than that later in training.
|
|
#self._update_param_scales(group, p, state, P_proj)
|
|
#self._update_param_scales_simple(group, p, state, P_proj)
|
|
|
|
# We won't be doing this any more.
|
|
#self._diagonalize_grad_cov(group, p, state)
|
|
self._zero_exp_avg_sq(state)
|
|
if step % grad_cov_period == 0:
|
|
self._update_grad_cov(group, p, state)
|
|
self._step(group, p, state)
|
|
|
|
state["step"] = step + 1
|
|
|
|
|
|
def _size_update(self,
|
|
group: dict,
|
|
scale_grads: Tensor,
|
|
p: Tensor,
|
|
state: dict) -> None:
|
|
"""
|
|
Called only where p.numel() > 1, this updates the scale of the parameter.
|
|
If we imagine: p = underlying_param * scale.exp(), and we are doing
|
|
gradient descent on underlying param and on scale, this function does the update
|
|
on `scale`.
|
|
|
|
Args:
|
|
group: dict to look up configuration values
|
|
scale_grads: a tensor of shape (size_update_period, batch_size, 1, 1,...) containing
|
|
grads w.r.t. the scales.
|
|
p: The parameter to update
|
|
state: The state-dict of p
|
|
"""
|
|
|
|
param_rms = state["param_rms"]
|
|
beta1, beta2 = group["betas"]
|
|
size_lr = group["lr"] * group["size_lr_scale"]
|
|
param_min_rms = group["param_min_rms"]
|
|
param_max_rms = group["param_max_rms"]
|
|
step = state["step"]
|
|
|
|
batch_size = p.shape[0]
|
|
size_update_period = scale_grads.shape[0]
|
|
# correct beta2 for the size update period: we will have
|
|
# faster decay at this level.
|
|
beta2_corr = beta2 ** size_update_period
|
|
|
|
scale_exp_avg_sq = state["scale_exp_avg_sq"] # shape: (batch_size, 1, 1, ..)
|
|
scale_exp_avg_sq.mul_(beta2_corr).add_(
|
|
(scale_grads ** 2).mean(dim=0), # mean over dim `size_update_period`
|
|
alpha=1-beta2_corr) # shape is (batch_size, 1, 1, ...)
|
|
|
|
# The 1st time we reach here is when size_step == 1.
|
|
size_step = (step + 1) // size_update_period
|
|
bias_correction2 = 1 - beta2_corr ** size_step
|
|
# we don't bother with bias_correction1; this will help prevent divergence
|
|
# at the start of training.
|
|
|
|
eps = 1.0e-10 # this value is not critical.
|
|
denom = scale_exp_avg_sq.sqrt() + eps # shape: (batch_size, 1, 1, ..)
|
|
|
|
scale_step = -size_lr * (bias_correction2 ** 0.5) * scale_grads.sum(dim=0) / denom
|
|
|
|
is_too_small = (param_rms < param_min_rms)
|
|
is_too_large = (param_rms > param_max_rms)
|
|
|
|
scale_step.masked_fill_(is_too_small, size_lr * size_update_period)
|
|
scale_step.masked_fill_(is_too_large, -size_lr * size_update_period)
|
|
delta = state["delta"]
|
|
# the factor of (1-beta1) relates to momentum.
|
|
delta.add_(p * scale_step, alpha=(1-beta1))
|
|
|
|
def _update_param_cov(self,
|
|
group: dict,
|
|
p: Tensor,
|
|
state: dict) -> None:
|
|
"""
|
|
Updates state[f"param_cov_{dim}"] for each dim of tensor p
|
|
(except batch and trivial and rank-1 dims)
|
|
"""
|
|
eps = group["eps"]
|
|
|
|
# zero_step is always the last time we called _update_param_cov.
|
|
# Our aim is to compute the parameter covariance averaged over all time
|
|
# (in steps) from the start, so "this_weight" that we give to this
|
|
# new covariance is proportional to the interval of time it represents,
|
|
# i.e. the interval from zero_step until step, while the existing "param_cov"
|
|
# represents the interval from step until zero_step.
|
|
# The min of 0.5 is a special case to ensure that the "torch.eye" we initialized
|
|
# param_cov with gets some weight on the 1st time we call this.
|
|
zero_step = state["zero_step"]
|
|
step = state["step"]
|
|
|
|
this_weight = min(0.5, (step - zero_step) / step)
|
|
|
|
batch_size = p.shape[0]
|
|
numel = p.numel() // batch_size
|
|
for dim in range(1, p.ndim):
|
|
size = p.shape[dim]
|
|
try:
|
|
param_cov = state[f"param_cov_{dim}"]
|
|
except KeyError:
|
|
continue # e.g. size == 1 or size == numel
|
|
|
|
# param_cov shape: (batch_size, num_blocks, block_size, block_size)
|
|
(batch_size, num_blocks, block_size, block_size) = param_cov.shape
|
|
|
|
# M: (batch_size, num_blocks, x, block_size)
|
|
M = p.transpose(dim, -1).reshape(batch_size, -1,
|
|
num_blocks, block_size).transpose(1, 2)
|
|
|
|
# will be batched matrix multiply. shape of this_param_cov:
|
|
# (batch_size, num_blocks, block_size, block_size)
|
|
this_param_cov = torch.matmul(M.transpose(2, 3), M)
|
|
|
|
# normalize scale of this param_cov, in case parameter scale
|
|
# changes significantly during training, which would cause some
|
|
# parts of the training timeline to be more highly weighted.
|
|
# shape of expression on r.h.s of "/=" has shape (batch_size, 1, 1, 1)
|
|
this_param_cov /= _mean(_diag(this_param_cov),
|
|
exclude_dims=[0], keepdim=True).unsqueeze(-1) + eps
|
|
|
|
param_cov.mul_(1-this_weight).add_(this_param_cov,
|
|
alpha=this_weight)
|
|
|
|
|
|
def _zero_exp_avg_sq(self,
|
|
state: dict) -> None:
|
|
"""
|
|
Zero the exp_avg_sq stats, and set state["zero_step"] to state["step"]
|
|
"""
|
|
state["exp_avg_sq"].zero_()
|
|
state["zero_step"] = state["step"]
|
|
|
|
def _is_lr_update_step(self,
|
|
group: dict,
|
|
state: dict) -> bool:
|
|
"""
|
|
Returns True if on this step we need to update the learning-rate matrices for this tensor
|
|
and False if not (note: this may just mean we update the gradient-dependent part).
|
|
The periodicity with which we update them increases from
|
|
(by default) 200 at the start of training to 2000 later on.
|
|
"""
|
|
try:
|
|
zero_step = state["zero_step"]
|
|
except:
|
|
# This parameter tensor has no learning-rate matrices to estimate
|
|
return False
|
|
step = state["step"]
|
|
period_initial, period_final = group["lr_update_period"]
|
|
# this formula gradually increases the periodicity from period_initial at the start
|
|
# to period_final when step >> 4 * period_final
|
|
cur_update_period = (period_initial +
|
|
((period_final - period_initial) * step /
|
|
(step + 4 * period_final)))
|
|
return (step >= zero_step + cur_update_period)
|
|
|
|
|
|
def _update_param_scales_simple(self,
|
|
group: dict,
|
|
p: Tensor,
|
|
state: dict,
|
|
P_proj: List[Optional[Tensor]]) -> None:
|
|
for dim in range(1, p.ndim):
|
|
size = p.shape[dim]
|
|
try:
|
|
Q = state[f"Q_{dim}"]
|
|
except KeyError:
|
|
assert size == 1 or size == numel, size
|
|
continue # e.g. size == 1 or size == numel
|
|
|
|
(batch_size, num_blocks, block_size, block_size) = Q.shape
|
|
this_P_proj = P_proj[dim].reshape(batch_size, num_blocks, block_size, 1)
|
|
|
|
# The following normalization step will ensure the Frobenius
|
|
# norm is unchanged, from applying this scale: at least,
|
|
# assuming "grad / denom" gives uncorrelated outputs so that
|
|
# they will have equal variances after projecting to the space
|
|
# where the parameter var is diagonalized... this is *roughly*
|
|
# true because the gradients at the point where we compute "grad
|
|
# / denom" should be decorrelated at least considering
|
|
# individual tensor dims
|
|
this_P_proj /= _mean(this_P_proj, exclude_dims=[0], keepdim=True)
|
|
|
|
if True:
|
|
# debug info.
|
|
scale = this_P_proj.sqrt()
|
|
step = state["step"]
|
|
scale_min, scale_max, scale_mean = scale.min().item(), scale.max().item(), scale.mean().item()
|
|
logging.info(f"step={step}, dim={dim}, size={size}, scale min,max,mean={scale_min,scale_max,scale_mean}")
|
|
|
|
Q *= this_P_proj.sqrt()
|
|
logging.info(f"Q rms = {(Q**2).mean().sqrt()}, abs-rms = {Q.abs().mean()}")
|
|
|
|
def _update_param_scales(self,
|
|
group: dict,
|
|
p: Tensor,
|
|
state: dict,
|
|
P_proj: List[Optional[Tensor]]) -> None:
|
|
"""
|
|
Modifies the scales on the rows of the learning-rate matrices Q for each dim of this tensor,
|
|
to take into account the estimated parameter covariance.
|
|
|
|
Args:
|
|
group: dict to look up configuration values
|
|
p: parameter matrix that we are updating. The learning rate matrices
|
|
are actually factors of p, so p itself will change when we change
|
|
them.
|
|
state: state dict for the current parameter
|
|
P_proj: a list indexed by dim, containing tensors of shape (batch_size, size)
|
|
where size == p.shape[dim], which represent the parameter covariance
|
|
projected by state[f"Q_{dim}"] (i.e. by the value of this variable
|
|
at entry to this function)
|
|
"""
|
|
ndim = p.ndim
|
|
batch_size = p.shape[0]
|
|
numel = p.numel() // batch_size
|
|
|
|
if numel in p[0].shape:
|
|
return # Nothing to do for this parameter matrix. E.g. a bias or a scalar.
|
|
|
|
scale_arr = [None] * ndim
|
|
|
|
# the small random part is to ensure there are no exact zeros, e.g.
|
|
# if we initialized some parameters to zero.
|
|
#
|
|
# we are going to make "cur_p" a corrected version of the parameter
|
|
# covariance that is projected with the orthogonal projections U, and that,
|
|
# on the axes given by these projections, has covariance proportional
|
|
# to that of state[f"param_cov_{dim}"] on those same axes. So it's
|
|
# *slightly* whitened, just to match the stats, but not completely
|
|
# whitened.
|
|
eps = 1.0e-20
|
|
cur_p = p + eps * torch.randn_like(p)
|
|
|
|
for dim in range(1, ndim):
|
|
size = p.shape[dim]
|
|
try:
|
|
Q = state[f"Q_{dim}"]
|
|
except KeyError:
|
|
assert size == 1 or size == numel, size
|
|
continue # e.g. size == 1 or size == numel
|
|
|
|
(batch_size, num_blocks, block_size, block_size) = Q.shape
|
|
|
|
M = cur_p.transpose(dim, -1)
|
|
# if p were of shape (batch_size, x, size, y, z),
|
|
# after the next line M would be of shape
|
|
# (batch_size, x, z, y, num_blocks, block_size)
|
|
M = M.reshape(*M.shape[:-1], num_blocks, block_size)
|
|
M = _move_dim(M, -2, 1) # (batch_size, num_blocks, x, z, y, block_size)
|
|
|
|
while Q.ndim < M.ndim:
|
|
Q = Q.unsqueeze(2)
|
|
# Now Q is of shape (batch_size, num_blocks, 1, 1, block_size, block_size)
|
|
# with indexes [batch_index, block_index, 1, 1, diagonalized_coordinate, canonical_coordinate],
|
|
# so we need to transpose Q as we convert M to the diagonalized co-ordinate.
|
|
M = torch.matmul(M, Q.transpose(-2, -1)) # (batch_size, num_blocks, x, z, y, block_size)
|
|
M = _move_dim(M, 1, -2) # (batch_size, x, z, y, num_blocks, block_size)
|
|
M = M.reshape(*M.shape[:-2], size) # (batch_size, x, z, y, size)
|
|
cur_p = M.transpose(dim, -1) # (batch_size, x, size, y, z)
|
|
|
|
# cur_param_var is a diagonal parameter variance over dimension `dim`,
|
|
# of the current "slightly-whitened" parameter; it
|
|
# will have shape [batch_size, 1, size, 1].
|
|
cur_param_var = _mean(cur_p**2,
|
|
exclude_dims=[0,dim],
|
|
keepdim=True) # (batch_size, 1, size, 1, 1) if dim==2
|
|
smoothed_param_var = P_proj[dim] # (batch_size, size)
|
|
S = smoothed_param_var.reshape(cur_param_var.shape) # (batch_size, 1, size, 1, 1) if dim==2
|
|
|
|
# OK, cur_param_var would have the values as S if the variance stats
|
|
# P_proj[dim] were accumulated from this exact parameter matrix and
|
|
# not smoothed, but actually they contain older versions of the
|
|
# parameter covariance and they have been smoothed, so they will, in
|
|
# general, be less extreme ("whiter spectrum"). We scale p so that
|
|
# it matches the estimated variance P_proj[dim]; the idea is to ensure it doesn't
|
|
# have too-extreme eigenvalues (where the stats permit).
|
|
# Actually we could just use P_proj[dim].sqrt(), suitably scaled,
|
|
# as the scales on the rows of Q (see _update_param_scales_simple() which does
|
|
# exactly this), but there is a problem of "counting things twice"
|
|
# which is easiest to understand for a 2-dimensional tensor, where the
|
|
# singular values show up identically in the covariance over either axis.
|
|
# The estimation procedure in this function avoids the "counting things twice"
|
|
# problem, at the expense of quite a bit of extra complexity.
|
|
scale = (S / cur_param_var.clamp(min=eps)).sqrt()
|
|
|
|
if True:
|
|
S_tmp = S.reshape(batch_size, size)
|
|
cur_tmp = cur_param_var.reshape(batch_size, size)
|
|
scale_tmp = scale.reshape(batch_size, size)
|
|
skip = 10 if size > 40 else 1
|
|
logging.info(f"dim={dim}/{ndim}, cur_param_var={cur_tmp[0][::skip]}, S={S_tmp[0][::skip]}, scale={scale_tmp[0][::skip]}")
|
|
|
|
|
|
if random.random() < 0.01:
|
|
skip = 10 if size < 20 else 1
|
|
logging.info(f"shape={p.shape}, dim={dim}, scale={scale[0].flatten()[::skip]}, cur_param_var={cur_param_var[0].flatten()[::skip]}, S={S[0].flatten()[::skip]}")
|
|
|
|
# scale shape: (batch_size, 1, size, 1, 1)
|
|
cur_p *= scale
|
|
|
|
# OK, at this point we have a matrix cur_p that is (somewhat)
|
|
# diagonalized by orthogonal matrices in each non-trivial dim, that is
|
|
# also multiplied by scalars to match the accumulated covariance stats,
|
|
# i.e. "slightly whitened" (in general this will make a modest
|
|
# difference, making the eigenvalue distribution a bit flatter). Now we
|
|
# will work out the "scaling" part of the learning-rate matrices Q. We
|
|
# can't do this independenty for each dim, because there is a risk of
|
|
# "counting things twice". E.g. for a matrix with 2 dims, if we do SVD
|
|
# M = U S V^T, if we consider the covariance on the left and the right,
|
|
# S will be reflected in both covariances, so it doesn't make sense to
|
|
# correct for S twice. Not all parameter matrices have exactly 2 dims,
|
|
# and also we're dealing with accumulated parameter stats which makes
|
|
# things not quite so simple, so we don't want to just take the sqrt of
|
|
# S.
|
|
|
|
# cur_scales[dim] will be a 1-d tensor with shapes like (batch_size, 1, 1, size, 1),
|
|
# containing the scales on the learning-rate matrix for this dimension.
|
|
# we apply these scales to the parameter matrix before estimating the
|
|
# cur_scales for the other dims.
|
|
cur_scales = [None] * ndim
|
|
|
|
debug = (random.random() < 0.1)
|
|
for i in range(4): # for 4 iterations (this is quite arbitrary)
|
|
for dim in range(1, ndim):
|
|
size = p.shape[dim]
|
|
if not f"Q_{dim}" in state:
|
|
assert size == 1 or size == numel, size
|
|
continue
|
|
if cur_scales[dim] is not None:
|
|
# correct for the fact that we have already normalized this dim in cur_p,
|
|
# i.e. remove the scaling for *this* dim, while keeping the scaling for
|
|
# all other dims.
|
|
cur_p *= cur_scales[dim] # cur_scales shape: (batch_size, 1, size, 1, 1)
|
|
# rms shape: (batch_size, 1, size, 1, 1)
|
|
rms = _mean(cur_p**2, exclude_dims=[0,dim], keepdim=True).sqrt()
|
|
rank = numel // size
|
|
# TODO: consider more smoothing here???
|
|
cur_scales[dim] = rms
|
|
cur_p /= rms
|
|
|
|
if debug:
|
|
def _summarize(rms):
|
|
rms = rms[0] # get rid of batch dim by selecting one example
|
|
rms = rms.flatten()
|
|
return rms[::10] # subset one every ten items
|
|
logging.info(f"i={i} shape={tuple(p.shape)}, dim={dim}, rank={rank}, size={size}, rms={_summarize(rms)}")
|
|
|
|
# Apply the scales in `cur_scales` to Q for each dim; this reflects the
|
|
# parameter rms values in the parameter-diagonalized space, that we have
|
|
# estimated in the loop above.
|
|
#
|
|
# We normalize the scales in such a way that the Frobenius norm
|
|
# after projecting (grad / denom) with Q should be unchanged, i.e. the
|
|
# same as (grad / denom), which is equivalent to having rms=1.0 due
|
|
# to how denom is constructed. This simplifies the normalization of the overall
|
|
# scale of the parameter change: we just have to multiply by the learning
|
|
# rate and param_rms.
|
|
for dim in range(1, ndim):
|
|
if cur_scales[dim] is not None:
|
|
size = p.shape[dim]
|
|
Q = state[f"Q_{dim}"]
|
|
(batch_size, num_blocks, block_size, block_size) = Q.shape
|
|
|
|
scale = cur_scales[dim].reshape(batch_size, num_blocks, block_size, 1)
|
|
|
|
# Geometrically interpolate scale with P_proj[dim].sqrt()
|
|
P_proj_weight = 0.5
|
|
scale = ((scale ** (1-P_proj_weight)) *
|
|
(P_proj[dim].reshape(batch_size, num_blocks, block_size, 1) ** (P_proj_weight * 0.5)))
|
|
|
|
# The following normalization step will ensure the Frobenius
|
|
# norm is unchanged, from applying this scale: at least,
|
|
# assuming "grad / denom" gives uncorrelated outputs so that
|
|
# they will have equal variances after projecting to the space
|
|
# where the parameter var is diagonalized... this is *roughly*
|
|
# true because the gradients at the point where we compute "grad
|
|
# / denom" should be decorrelated at least considering
|
|
# individual tensor dims
|
|
scale /= _mean(scale**2, exclude_dims=[0], keepdim=True).sqrt()
|
|
|
|
|
|
if True:
|
|
# debug info.
|
|
step = state["step"]
|
|
scale_min, scale_max, scale_mean = scale.min().item(), scale.max().item(), scale.mean().item()
|
|
logging.info(f"step={step}, dim={dim}, size={size}, scale min,max,mean={scale_min,scale_max,scale_mean}")
|
|
|
|
# Q is indexed [batch_index, block_index, diagonalized_coordinate, canonical_coordinate],
|
|
# want to multiply on the diagonalized co-ordinate.
|
|
state[f"Q_{dim}"] *= scale
|
|
state["last_param_scale_update"] = state["step"]
|
|
|
|
|
|
def _compute_bases(self,
|
|
group: dict,
|
|
p_shape: torch.Size,
|
|
state: dict) -> List[Optional[Tensor]]:
|
|
"""
|
|
For each tensor dimension that we are preconditioning, this function sets
|
|
state[f"Q_{dim}"] to an orthogonal matrix (per batch element and
|
|
block); it will be indexed [batch_index, block_index, diagonalized_coordinate, canonical_coordinate].
|
|
This will be the matrix that diagonalizes Z,
|
|
where Z is the symmetric matrix satisfying:
|
|
Z G Z = P (eqn:1)
|
|
with G being the gradient covariance grad_cov, P being the smoothed version of the parameter
|
|
covariance param_cov, and Z the symmetric positive definit learning-rate matrices.
|
|
We'll discuss later how we solve for Z. Firstly, because we want
|
|
a basis that is as meaningful possible for smoothing P, we will first put P and G in a basis
|
|
where G is diagonal. That is, we do an SVD G = U G' U^T, and multiplying (eqn:1) on the
|
|
left and right by U^T and U respetively, and inserting some expressions equivalent to I,
|
|
we have:
|
|
(U^T Z U) (U^T G U) (U^T Z U) = (U^T P U)
|
|
or: Z' G' Z' = P' (eqn:2)
|
|
where Z' = U^T Z U and P' = U^T P U, and of course G' is diagonal. We are actually going to
|
|
be solving (eqn:2), and then computing:
|
|
Z = U Z' U^T.
|
|
|
|
A solution to (eqn:2) is as follows. We are going to be using a Cholesky-based solution in
|
|
favor of one that requires SVD or eigenvalue decomposition, because it is much faster (we first
|
|
have to be careful that the input is not close to singular, though).
|
|
|
|
So, with the Cholesky decomposition of P' being:
|
|
P' = C C^T,
|
|
the solution is going to be:
|
|
Z' = C (C^T G' C)^{-0.5} C^T (eqn:3)
|
|
[[ we can verify that this is a solution by multiplying (eqn:1) by C^{-1} and C^{-T} on the
|
|
left and right respectively, giving:
|
|
(C^T G' C)^{-0.5} C^T G' C (C^T G' C)^{-0.5} = I
|
|
which we can immediately see is the case. ]]
|
|
|
|
Args:
|
|
group: dict to look up config values
|
|
p_shape: the shape of the batch of identical-sized tensors we are optimizing
|
|
state: dict to look up optimization state
|
|
|
|
Return:
|
|
This function returns a list of Tensors P_proj indexed by dim from 0..ndim-1, the tensors being
|
|
of shape (batch_size, size), which contain the diagonal of the smoothed parameter covariance
|
|
P after projection by the orthogonal matrix state[f"Q_{dim}"] that diagonalizes Z (this is
|
|
the space in which we do the actual update).
|
|
"""
|
|
ndim = len(p_shape)
|
|
P_proj = [None] * ndim
|
|
for dim in range(1, ndim):
|
|
size = p_shape[dim]
|
|
try:
|
|
Q = state[f"Q_{dim}"]
|
|
except KeyError:
|
|
assert size == 1 or size == numel, size
|
|
continue # e.g. size == 1 or size == numel
|
|
|
|
|
|
param_cov = state[f"param_cov_{dim}"] # (batch_size, num_blocks, block_size, block_size)
|
|
grad_cov = state[f"grad_cov_{dim}"] # (batch_size, num_blocks, block_size, block_size)
|
|
(batch_size, num_blocks, block_size, block_size) = param_cov.shape
|
|
# U_g diagonalizes grad_cov, in the sense that U_g^T grad_cov U_g is diagonal; its shape
|
|
# is the same as param_cov and grad_cov.
|
|
#
|
|
# G_prime is the diagonalized grad_cov, G' above, of shape (batch_size, num_blocks, block_size).
|
|
U_g, G_prime, _ = _svd(grad_cov)
|
|
# P_prime is P' above, which represents param_cov in the basis that diagonalizes G_prime.
|
|
# It is not smoothed yet.
|
|
P_prime = torch.matmul(U_g.transpose(2, 3), torch.matmul(param_cov, U_g))
|
|
|
|
P_prime_unsmoothed = P_prime
|
|
P_prime = self._smooth_param_cov(group, p_shape, P_prime, G_prime)
|
|
|
|
# C will satisfy: P_prime == torch.matmul(C, C.transpose(2, 3))
|
|
# C is of shape (batch_size, num_blocks, block_size, block_size).
|
|
#def _fake_cholesky(X):
|
|
# U_, S_, _ = _svd(X)
|
|
# return U_ * S_.sqrt().unsqueeze(-2)
|
|
#C = _fake_cholesky(P_prime)
|
|
C = P_prime.cholesky()
|
|
|
|
|
|
# CGC = (C^T G' C), it would be torch.matmul(C.transpose(2, 3), torch.matmul(G, C))
|
|
# if G were formatted as a diagonal matrix, but G is just the diagonal.
|
|
CGC = torch.matmul(C.transpose(2, 3) * G_prime.unsqueeze(-2), C)
|
|
|
|
U, S, _ = _svd(CGC) # Need SVD to compute CGC^{-0.5}
|
|
# next we compute (eqn:3). The thing in the parenthesis is, CGC^{-0.5},
|
|
# can be written as U S^{-0.5} U^T, so the whole thing is
|
|
# (C U) S^{-0.5} (C U)^T.
|
|
#
|
|
# (C U) has good condition number because we smoothed the
|
|
# eigenvalues of P_prime; S^{-0.5} may not have good condition, in
|
|
# fact S could have zero eigenvalues (as a result of G having zero
|
|
# or very small elements). For numerical reasons, to avoid infinity
|
|
# in elements of S when we compute S^{-0.5}, we apply a floor here.
|
|
# Later, we will only care about the eigen-directions of Z_prime and
|
|
# not the eigenvalues, so aside from numerical consideratins we
|
|
# don't want to smooth S too aggressively here.
|
|
eps = 1.0e-12
|
|
S_floor = eps + (1.0e-06 * _mean(S, exclude_dims=[0], keepdim=True))
|
|
|
|
S = _soft_floor(S, S_floor)
|
|
CU = torch.matmul(C, U)
|
|
S_inv_sqrt = 1.0 / S.sqrt()
|
|
Z_prime = torch.matmul(CU * S_inv_sqrt.unsqueeze(-2),
|
|
CU.transpose(2, 3))
|
|
|
|
if True:
|
|
def _check_similar(x, y, name):
|
|
ratio = (y-x).abs().sum() / x.abs().sum()
|
|
if not (ratio < 0.0001):
|
|
logging.warn(f"Check {name} failed, ratio={ratio.item()}")
|
|
|
|
def _check_symmetric(x, x_name):
|
|
diff = x - x.transpose(-2, -1)
|
|
ratio = diff.abs().sum() / x.abs().sum()
|
|
if not (ratio < 0.0001):
|
|
logging.warn(f"{x_name} is not symmetric: ratio={ratio.item()}")
|
|
|
|
_check_similar(P_prime, torch.matmul(C, C.transpose(2, 3)), "cholesky_check")
|
|
_check_symmetric(Z_prime, "Z_prime")
|
|
_check_symmetric(P_prime, "P_prime")
|
|
# A check.
|
|
# Z_prime is supposed to satisfy: Z_prime G_prime Z_prime = P_prime (eqn:2)
|
|
P_prime_check = torch.matmul(Z_prime * G_prime.unsqueeze(-2), Z_prime)
|
|
diff_ratio_l2 = ((P_prime - P_prime_check) ** 2).sum().sqrt() // (P_prime ** 2).sum().sqrt()
|
|
diff_ratio_l1 = (P_prime - P_prime_check).abs().sum() // P_prime.abs().sum()
|
|
if not (diff_ratio_l2 < 0.00001) or diff_ratio_l1 > 0.03:
|
|
logging.warn(f"Z_prime does not satisfy its definition, diff_ratio_l{1,2} = {diff_ratio_l1.item(),diff_ratio_l2.item()}, size={size}")
|
|
|
|
|
|
# We really want the SVD on Z, which will be used for the learning-rate matrix
|
|
# Q, but Z_prime is better, numerically, to work on because it's closer to
|
|
# being diagonalized.
|
|
U_z_prime, S, _ = _svd(Z_prime)
|
|
|
|
U_z = torch.matmul(U_g, U_z_prime)
|
|
# We could obtain Z in two possible ways.
|
|
# Commenting the check below.
|
|
# Z_a = torch.matmul(U_z * S.unsqueeze(-2), U_z.transpose(2, 3))
|
|
# Z_b = torch.matmul(U_g, torch.matmul(Z_prime, U_g.transpose(2, 3)))
|
|
# _check_similar(Z_a, Z_b, "Z")
|
|
## OK, Z is the SPD transform that maps G to P, as in Z G Z = P.
|
|
## We just need the basis U_z that diagonalizes this.
|
|
## U_z, S, _ = _svd(Z)
|
|
if True:
|
|
skip = 10 if S.shape[-1] > 40 else 1
|
|
logging.info(f"dim={dim}, G_prime is {G_prime[0,0,::skip]}, Eigs of Z are: {S[0,0,::skip]}")
|
|
|
|
# state[f"Q_{dim}"] is indexed: [batch_idx, block_idx, diagonalized_coordinate, canonical_coordinate].
|
|
# so we need to transpose U_z as U_z is indexed
|
|
# [batch_idx, block_idx, canonical_coordinate, diagonalized_coordinate]
|
|
Q[:] = U_z.transpose(2, 3)
|
|
|
|
# Work out the projected smoothed parameter covariance P_proj, which is P
|
|
# projected in the basis U_z. Now,
|
|
# P = U_g P' U_g^T,
|
|
# and P_proj = U_z^T P U_z,
|
|
# so P_proj = (U_z^T U_g) P' (U_z^T U_g)^T
|
|
U_prod = torch.matmul(U_z.transpose(2, 3), U_g)
|
|
# this_P_proj shape: (batch_size, num_blocks, block_size)
|
|
this_P_proj = _diag(torch.matmul(U_prod, torch.matmul(P_prime, U_prod.transpose(2, 3))))
|
|
P_proj[dim] = this_P_proj.clone().reshape(batch_size, size)
|
|
|
|
simple_update = True
|
|
if simple_update:
|
|
# normalize the scales in a way that preserves the Frobenius norm of the
|
|
# projected parameter deltas
|
|
P_rms = this_P_proj / _mean(this_P_proj, exclude_dims=[0], keepdim=True)
|
|
scale = P_rms.unsqueeze(-1).sqrt()
|
|
Q *= scale
|
|
logging.info(f"Q rms = {(Q**2).mean().sqrt()} abs-rms = {Q.abs().mean()}")
|
|
# no iterative stuff, just use sqrt(P_proj) as scale on Q. If this is False, we need to
|
|
# call self._update_param_scales(...) from the calling function.
|
|
if True:
|
|
# debug output
|
|
step = state["step"]
|
|
scale_min, scale_max, scale_mean = scale.min().item(), scale.max().item(), scale.mean().item()
|
|
logging.info(f"step={step}, dim={dim}, size={size}, scale min,max,mean={scale_min,scale_max,scale_mean}")
|
|
if True:
|
|
# debug output
|
|
this_P_proj_unsmoothed = _diag(torch.matmul(U_prod, torch.matmul(P_prime_unsmoothed,
|
|
U_prod.transpose(2, 3))))
|
|
G_proj_unsmoothed = _diag(torch.matmul(U_prod * G_prime.unsqueeze(-2), U_prod.transpose(2, 3)))
|
|
skip = 10 if block_size > 40 else 1
|
|
logging.info(f"dim={dim}, diag of P_proj is: {P_proj[dim][0,0:block_size:skip]}, diag of unsmoothed P_proj is: {this_P_proj_unsmoothed[0,0,::skip]}, diag of unsmoothed G_proj is {G_proj_unsmoothed[0,0,::skip]}")
|
|
|
|
# P_proj won't be needed if simple_update == True.
|
|
return P_proj
|
|
|
|
|
|
def _smooth_param_cov(self,
|
|
group: dict,
|
|
p_shape: torch.Size,
|
|
P_prime: Tensor,
|
|
G_prime: Tensor) -> Tensor:
|
|
"""
|
|
This function returns a modified/smoothed version of the parameter covariance
|
|
P_prime.
|
|
|
|
Args:
|
|
group: dict to look up config values
|
|
p_shape: The shape of the parameter we are optimizing
|
|
P_prime: a Tensor of shape (batch_size, num_blocks, block_size, block_size),
|
|
containing the parameter covariance in a basis that diagonalizes the
|
|
gradient covariance.
|
|
G_prime: the diagonalized gradient covariance, of shape (batch_size, num_blocks,
|
|
block_size)
|
|
|
|
|
|
state[f"param_cov_{dim}"], which is an estimate of the covariance of the parameter
|
|
p, averaged over time, and taken over dimension `dim` of the tensor.
|
|
|
|
The smoothing done here limits the extent to which the parameter covariance
|
|
can be strongly "off-diagonal" with respect to the gradient covariance. That is:
|
|
if the parameter covariance is just the gradient covariance to some power, this
|
|
function does no smoothing; but if it is highly off-diagonal we do more smoothing.
|
|
"""
|
|
P_prime_diag = _diag(P_prime) # (batch_size, num_blocks, block_size)
|
|
eps = 1.0e-10
|
|
P_prime_diag = (P_prime_diag + eps) / P_prime_diag.mean()
|
|
# make sure no diagonal element is close to zero.. we don't expect this
|
|
# would happen. this is likely not important. Note, this just used for
|
|
# normalizing P prior to smoothing.
|
|
P_prime_diag.clamp_(min=0.01)
|
|
P_prime_rms = P_prime_diag.sqrt()
|
|
P_prime_scale = P_prime_rms.unsqueeze(-1) * P_prime_rms.unsqueeze(-2)
|
|
|
|
# P_norm will have diagonal elements close to 1. We do some smoothing
|
|
# in this space.
|
|
P_norm = P_prime / P_prime_scale
|
|
# Now P is as normalized as we can make it... do smoothing baserd on 'rank',
|
|
# that is intended to compensate for bad estimates of P.
|
|
(batch_size, num_blocks, block_size, block_size) = P_norm.shape
|
|
# `rank_per_block` is the rank of each block of P_prime if we were to estimate it from just one
|
|
# parameter tensor. We average it over time, but actually it won't be changing
|
|
# too much, so `rank` does tell us something.
|
|
size = num_blocks * block_size
|
|
rank = p_shape.numel() // (size * batch_size) # actually the rank of each block
|
|
smooth0 = group["param_rms_smooth0"]
|
|
smooth1 = group["param_rms_smooth1"]
|
|
# We want expr for smoothing amount to be of the form: smooth = alpha * size / (beta*rank + size)
|
|
# for "size" here, we actually want to use block_size, since we are concerned about the
|
|
# robustness of the covariance within these blocks.
|
|
# param_rms_smooth{0,1} represents the user-specified desired amount of smoothing
|
|
# when rank==0*size and rank==1*size, respectively.
|
|
# from rank==0*size, we get smooth0 = alpha * size/size, so alpha = smooth0.
|
|
# from setting rank==size, we get smooth1 = alpha * size / (beta*size * size) = alpha/(1+beta),
|
|
# so smooth1 == smooth0 / (1+beta), so (1+beta) = smooth0/smooth1, so beta=smooth0/smooth1 - 1
|
|
smooth = smooth0 * block_size / ((smooth0/smooth1 - 1) * rank + block_size)
|
|
if True:
|
|
logging.info(f"block size={block_size}, rank={rank}, smooth={smooth}")
|
|
# add rank-dependent smoothing amount to diagonal of P_prime. _diag() returns an aliased tensor.
|
|
# we don't need to multiply `smooth` by anything, because at this point, P_prime should have
|
|
# diagonal elements close to 1.
|
|
_diag(P_norm).add_(smooth)
|
|
|
|
P_norm = self._smooth_cov(P_norm,
|
|
group["cov_min"][0],
|
|
group["cov_max"][0],
|
|
group["cov_pow"][0])
|
|
# Remove the diagonal preconditioning on P_norm, giving us stage-1-smoothed
|
|
# version of P_prime.
|
|
P_prime = P_norm * P_prime_scale
|
|
|
|
if True:
|
|
# This block smooths G_prime.
|
|
# Make sure G_prime has unit mean and no eigenvalue is super small. Note, G_prime
|
|
# is already diagonalized, the variable G_prime is just the tensor of eigenvalues.
|
|
G_prime_mean = _mean(G_prime, exclude_dims=[0], keepdim=True)
|
|
G_prime_min = group["cov_min"][3]
|
|
# make sure G_prime has no zero eigs, and is unit mean.
|
|
#G_prime = _soft_floor(G_prime, G_prime_min * G_prime_mean, depth=3)
|
|
G_prime = G_prime + G_prime_min * G_prime_mean + 1.0e-20
|
|
G_prime /= _mean(G_prime, exclude_dims=[0], keepdim=True)
|
|
# it now has unit mean..
|
|
G_prime_max = group["cov_max"][3]
|
|
G_prime = 1. / (1./G_prime + 1./G_prime_max) # apply max
|
|
G_prime_pow = group["cov_pow"][3]
|
|
if G_prime_pow != 1.0:
|
|
G_prime = G_prime ** G_prime_pow
|
|
|
|
G_prime_rms = G_prime.sqrt()
|
|
G_prime_scale = G_prime_rms.unsqueeze(-1) * G_prime_rms.unsqueeze(-2)
|
|
# P_gnorm is a version of P_prime that is multiplied by G (actually
|
|
# G^{0.5} P_prime G^{0.5}), so that it reflects the amount of
|
|
# loss-function change in each dimension. We also tried smoothing
|
|
# a version of P_prime divided by G, but it seemed not to be helpful.
|
|
P_gnorm = P_prime * G_prime_scale
|
|
# Apply another round of smoothing "relative to G"
|
|
P_gnorm = self._smooth_cov(P_gnorm,
|
|
group["cov_min"][1],
|
|
group["cov_max"][1],
|
|
group["cov_pow"][1])
|
|
# Undo the scaling relative to G, so we have stage-2-smoothed version of P_prime.
|
|
P_prime = P_gnorm / G_prime_scale
|
|
|
|
# Apply a 3rd round of smoothing in the canonical basis.
|
|
P_prime = self._smooth_cov(P_prime,
|
|
group["cov_min"][2],
|
|
group["cov_max"][2],
|
|
group["cov_pow"][2])
|
|
return P_prime
|
|
|
|
def _smooth_cov(self,
|
|
X: Tensor,
|
|
min_eig: float,
|
|
max_eig: float,
|
|
power: float = 1.0) -> Tensor:
|
|
"""
|
|
Returns a `smoothed` version of a symmetric positive definite covariance matrix
|
|
[with block-diagonal structure, in a batch]. This is done without SVD (which
|
|
can be very slow).
|
|
The eigenvalues L will be transformed as:
|
|
|
|
L = L + min_eig * L.mean() + eps
|
|
L /= L.mean()
|
|
L = 1 / (1/L + 1/max_eig) # soft-min between L and max_eig
|
|
L = L ** power # need SVD for this, will get rid of the requirement later.
|
|
L /= L.mean()
|
|
|
|
# Note on approximation functions like x^0.75 for smallish x: on wolframalpha, type:
|
|
# plot x^0.75 and 0.05 + (1.1x - 0.18 x^2 + 0.02 x^3) for x from 0 to 10
|
|
# [this starts to diverge after 5 or so]
|
|
|
|
Args:
|
|
min_eig: minimum allowed eigenvalue of returned X
|
|
max_eig: maximum allowed eigenvalue of returned X
|
|
power: power to take eigenvalues to
|
|
X: the batch of symmetric positive definite tensors we are smoothing;
|
|
of shape (batch_size, num_blocks, block_size, block_size)
|
|
"""
|
|
eps = 1.0e-10
|
|
if power != 1.0:
|
|
U, S, _ = _svd(X)
|
|
S_mean = _mean(S, exclude_dims=[0], keepdim=True)
|
|
S = S + min_eig * S_mean + eps
|
|
S_mean = S_mean * (1 + min_eig) + eps
|
|
S = S / S_mean
|
|
S = 1. / (1./S + 1./max_eig)
|
|
S = S ** power
|
|
S = S / _mean(S, exclude_dims=[0], keepdim=True)
|
|
return torch.matmul(U * S.unsqueeze(-2), U.transpose(2, 3))
|
|
else:
|
|
X = X.clone()
|
|
diag = _diag(X) # Aliased with X
|
|
mean_eig = _mean(diag, exclude_dims=[0], keepdim=True)
|
|
diag += (mean_eig * min_eig + eps)
|
|
cur_diag_mean = mean_eig * (1 + min_eig) + eps
|
|
# The following 2 statements will be equivalent to:
|
|
# L /= L.mean()
|
|
# L = 1 / (1/L + 1/max_eig) # soft-min between L and max_eig
|
|
# if L is the eigenvalues
|
|
X_inv = X.inverse()
|
|
_diag(X_inv).add_(1. / (max_eig * cur_diag_mean))
|
|
X = X_inv.inverse()
|
|
X /= _mean(_diag(X), exclude_dims=[0], keepdim=True).unsqueeze(-1)
|
|
return X
|
|
|
|
|
|
def _diagonalize_grad_cov(self,
|
|
group: dict,
|
|
p: Tensor,
|
|
state: dict) -> None:
|
|
"""
|
|
Called from _update_lrs(), this function diagonalizes the gradient covariance
|
|
state[f"grad_cov_{dim}"] by modifying the projections state[f"Q_{dim}"]: specifically,
|
|
by left-multiplying the projections by an orthogonal matrix.
|
|
"""
|
|
batch_size = p.shape[0]
|
|
numel = p.numel() // batch_size
|
|
for dim in range(1, p.ndim): # dim 0 is batch dim
|
|
size = p.shape[dim]
|
|
try:
|
|
# A and grad_cov shape: (batch_size, num_blocks, block_size, block_size)
|
|
Q = state[f"Q_{dim}"]
|
|
grad_cov = state[f"grad_cov_{dim}"]
|
|
except KeyError:
|
|
assert size == 1 or size == numel
|
|
continue
|
|
|
|
# Suppose the actual parameter matrix p is M, of shape (-1, size), where
|
|
# the -1 represents all other tensor dims treated as a batch dimension.
|
|
# M_grad is the same shape as M. We could write a pseudo-loss as
|
|
# loss = tr(M_grad^T M)
|
|
# Because we can decompose M as M == N Q, we can write:
|
|
# loss = tr(M_grad^T N Q) = tr(Q M_grad^T N),
|
|
# so we can write this at tr(N_grad^T N),
|
|
# where N_grad == (Q M_grad^T)^T = M_grad Q^T.
|
|
# Now,
|
|
# grad_cov == M_grad^T M_grad,
|
|
# decaying-averaged over minibatches; this is of shape (size,size).
|
|
# Using N_grad = M_grad Q^T, we can write the gradient covariance w.r.t. N
|
|
# (which is what we want to diagonalize), as::
|
|
# N_grad_cov = N_grad^T N_grad
|
|
# = Q M_grad^T M_grad Q^T
|
|
# = Q grad_cov Q^T
|
|
# (note: this makes sense because the 1st index of Q is the diagonalized
|
|
# index).
|
|
|
|
def dispersion(v):
|
|
# a ratio that, if the eigenvalues are more dispersed, it will be larger.
|
|
return '%.3e' % ((v**2).mean() / (v.mean() ** 2)).item()
|
|
|
|
|
|
# N_grad_cov shape: (batch_size, num_blocks, block_size, block_size)
|
|
N_grad_cov = torch.matmul(Q, torch.matmul(grad_cov, Q.transpose(2, 3)))
|
|
N_grad_cov = N_grad_cov + N_grad_cov.transpose(2, 3) # ensure symmetric
|
|
U, S, V = _svd(N_grad_cov)
|
|
if random.random() < 0.01:
|
|
logging.info(f"Diagonalizing, shape={tuple(p.shape)}, dim={dim}, dispersion "
|
|
f"changed from {dispersion(_diag(N_grad_cov))} to {dispersion(S)}")
|
|
|
|
# N_grad_cov is SPD, so
|
|
# N_grad_cov = U S U^T.
|
|
|
|
# Now, we can diagonalize N_grad_cov with:
|
|
# U^T N_grad_cov U == S.
|
|
# N_grad_cov is a sum of N_grad^T N_grad.
|
|
# We know U^T N_grad_cov U is diagonal, so U^T N_grad^T N_grad U is diagonal.
|
|
# The linearized pseudo-loss can be written as tr(N_grad^T N_grad).
|
|
# This can be written as tr(U U^T N_grad^T N_grad), since U U^T == I,
|
|
#
|
|
# which we can rearrange as tr(U^T N_grad^T N U). This can be interpreted
|
|
# as tr(hat_N_grad hat_N), where:
|
|
# hat_N_grad = N_grad U
|
|
# hat_N = N U
|
|
# (hat_N means \hat{N}, or N with a hat on it).
|
|
# So if we interpret hat_N = N U, the gradient covariance w.r.t.
|
|
# hat_N will be diagonalized. We also modify Q to hat_Q when
|
|
# we modify hat_N, to keep the product M unchanged:
|
|
# M = N Q = N U U^T Q = hat_N hat_Q
|
|
# This can be done by setting
|
|
# hat_Q := U^T Q (eq.10)
|
|
#
|
|
# This is the only thing we have to do, as N is implicit
|
|
# and not materialized at this point.
|
|
Q[:] = torch.matmul(U.transpose(2, 3), Q)
|
|
|
|
def _update_grad_cov(self,
|
|
group: dict,
|
|
p: Tensor,
|
|
state: dict):
|
|
"""
|
|
Updates the gradient covariance state[f"grad_cov_{dim}"] for appropriate
|
|
dimensions of the tensor.
|
|
"""
|
|
beta2 = group["betas"][1]
|
|
grad = p.grad
|
|
batch_size = p.shape[0]
|
|
for dim in range(1, grad.ndim): # dim 0 is batch dim.
|
|
size = grad.shape[dim]
|
|
name = f"grad_cov_{dim}"
|
|
if name not in state:
|
|
continue # e.g. size==1, size == numel, size > max_block_size
|
|
grad_cov = state[name] # shaped: (batch_size, num_blocks, block_size, block_size)
|
|
(batch_size, num_blocks, block_size, block_size) = grad_cov.shape
|
|
|
|
g = grad.transpose(-1, dim)
|
|
# if grad were of shape (batch_size, x, size, y, z),
|
|
# g would be of shape (batch_size, x, z, y, size); and
|
|
# after the next line, g will be of shape
|
|
# (batch_size, x, z, y, num_blocks, block_size)
|
|
g = g.reshape(*g.shape[:-1], num_blocks, block_size)
|
|
g = _move_dim(g, -2, 1) # (batch_size, num_blocks, x, z, y, block_size)
|
|
g = g.reshape(batch_size, num_blocks, -1, block_size)
|
|
# now g is of shape (batch_size, num_blocks, x*z*y, block_size)
|
|
|
|
# this_grad_cov: (batch_size, num_blocks, block_size, block_size)
|
|
this_grad_cov = torch.matmul(g.transpose(-2, -1), g)
|
|
|
|
grad_cov.mul_(beta2).add_(this_grad_cov, alpha=(1-beta2))
|
|
|
|
def _step(self,
|
|
group: dict,
|
|
p: Tensor,
|
|
state: dict):
|
|
"""
|
|
This function does the core update of self._step, in the case where the tensor
|
|
has more than 1 elements. It multiplies the moving-average gradient tensor by a
|
|
state["lr_{dim}"] on each non-trivial axis/dimension, does a length normalization,
|
|
and takes a step in that direction.
|
|
|
|
Args:
|
|
group: A dict which will be used to look up configuration values
|
|
p: The parameter to be updated
|
|
grad: The grad of p
|
|
state: The state-dict corresponding to parameter p
|
|
batch: if true, the first dimension of p is a batch dim, i.e. a batch of
|
|
parameter matrices.
|
|
|
|
This function modifies p.
|
|
"""
|
|
grad = p.grad
|
|
lr = group["lr"]
|
|
beta1, beta2 = group["betas"]
|
|
eps = group["eps"]
|
|
step = state["step"]
|
|
|
|
grad = self._project(grad, state, forward=True)
|
|
|
|
exp_avg_sq = state["exp_avg_sq"]
|
|
exp_avg_sq.mul_(beta2).addcmul_(grad, grad,
|
|
value=(1-beta2))
|
|
|
|
this_step = state["step"] - (state["zero_step"] if "zero_step" in state else 0)
|
|
bias_correction2 = 1 - beta2 ** (this_step + 1)
|
|
if bias_correction2 < 0.99:
|
|
# note: not in-place.
|
|
exp_avg_sq = exp_avg_sq * (1.0 / bias_correction2)
|
|
denom = exp_avg_sq.sqrt() + eps
|
|
|
|
grad = grad / denom
|
|
|
|
# and project back..
|
|
grad = self._project(grad, state, forward=False)
|
|
|
|
|
|
alpha = -lr * (1-beta1) * state["param_rms"]
|
|
|
|
delta = state["delta"]
|
|
delta.add_(grad * alpha)
|
|
p.add_(delta)
|
|
|
|
|
|
def _step_scalar(self,
|
|
group: dict,
|
|
p: Tensor,
|
|
state: dict):
|
|
"""
|
|
A simplified form of the core update for scalar tensors, where we cannot get a good
|
|
estimate of the parameter rms. p will not be a scalar, because it has a batch dim.
|
|
"""
|
|
beta1, beta2 = group["betas"]
|
|
scalar_max = group["scalar_max"]
|
|
eps = group["eps"]
|
|
lr = group["lr"]
|
|
grad = p.grad
|
|
|
|
exp_avg_sq = state["exp_avg_sq"] # shape: (batch_size,)
|
|
exp_avg_sq.mul_(beta2).addcmul_(grad, grad,
|
|
value=1-beta2)
|
|
|
|
# bias_correction2 is like in Adam. Don't bother with bias_correction1;
|
|
# slower update at the start will help stability anyway.
|
|
bias_correction2 = 1 - beta2 ** (state["step"] + 1)
|
|
denom = (exp_avg_sq / bias_correction2).sqrt() + eps
|
|
|
|
delta = state["delta"]
|
|
delta.add_(grad / denom, alpha=-lr*(1-beta1))
|
|
p.clamp_(min=-scalar_max, max=scalar_max)
|
|
p.add_(delta)
|
|
|
|
def _project(self,
|
|
M: Tensor,
|
|
state: dict,
|
|
forward: bool) -> Tensor:
|
|
"""
|
|
Multiply a tensor M by a projection Q on each of its dimensions (for which we have
|
|
such a projection)
|
|
|
|
Args:
|
|
M: The tensor to project, e.g. a gradient or a parameter change.
|
|
state: dict to look up the projections
|
|
forward: if True, go in the forward direction (from canonical to diagnonalized
|
|
co-ordinates); if False, the reverse. They differ by a transpose,
|
|
not an inverse, and do not make a round trip.
|
|
"""
|
|
numel = M.numel()
|
|
|
|
for dim in range(1, M.ndim): # dim 0 is batch dim (batch of parameter
|
|
# tensors of same shape)
|
|
try:
|
|
Q = state[f"Q_{dim}"] # shape: (batch_size, num_blocks, block_size, block_size)
|
|
except KeyError:
|
|
continue # no projection for this dim
|
|
|
|
(batch_size, num_blocks, block_size, block_size) = Q.shape
|
|
size = M.shape[dim] # == num_blocks * block_size
|
|
|
|
if forward:
|
|
# Q is indexed
|
|
# [batch_idx, block_idx, diagonalized_idx, canonical_idx]; in
|
|
# the forward direction we want to change canonical to
|
|
# diagonalized index, so we transpose.
|
|
Q = Q.transpose(2, 3)
|
|
|
|
# assume M currently has shape (batch_size, x, size, y, z), and dim==2.
|
|
M = M.transpose(-1, dim) # (batch_size, x, y, z, size)
|
|
M = M.reshape(batch_size, *M.shape[1:-1],
|
|
num_blocks, block_size) # (batch_size, x, y, z, num_blocks, block_size)
|
|
M = _move_dim(M, -2, 1) # (batch_size, num_blocks, x, y, z, block_size)
|
|
|
|
while Q.ndim < M.ndim:
|
|
Q = Q.unsqueeze(2)
|
|
# now Q has shape (batch_size, num_blocks, 1, 1, block_size, block_size)
|
|
if M.ndim < Q.ndim: # special case where M shape is e.g. (batch_size, num_blocks, x)
|
|
M = torch.matmul(M.unsqueeze(-2), Q).squeeze(-2)
|
|
else:
|
|
M = torch.matmul(M, Q) # (batch_size, num_blocks, x, y, z, block_size)
|
|
M = _move_dim(M, 1, -2) # (batch_size, x, y, z, num_blocks, block_size)
|
|
M = M.reshape(*M.shape[:-2], size) # (batch_size, x, y, z, size)
|
|
M = M.transpose(-1, dim) # (batch_size, x, size, y, z)
|
|
return M
|
|
|
|
|
|
def _move_dim(x: Tensor, orig_dim: int, new_dim: int) -> Tensor:
|
|
"""
|
|
Moves the position of one dimension in a Tensor, while keeping
|
|
the remaining dims in the same order. E.g. if x has
|
|
shape [10, 1, 2, 3], _move_dim(x, 3, 1) will have shape
|
|
[10, 3, 1, 2]. Returns the reshaped tensor which will be
|
|
a view of the original Tensor.
|
|
Args:
|
|
x: Tensor to reshape
|
|
orig_dim: Original dimension to move, with -x.ndim <= orig_dim < x.ndim
|
|
new_dim: New position of the original dimension, with
|
|
-x.ndim <= new_dim < x.ndim
|
|
Returns: a permuted view of x
|
|
"""
|
|
if orig_dim < 0:
|
|
orig_dim += x.ndim
|
|
if new_dim < 0:
|
|
new_dim += x.ndim
|
|
dims = list(range(x.ndim))
|
|
dims[orig_dim] = -1
|
|
dims.remove(-1)
|
|
dims.insert(new_dim, orig_dim)
|
|
return x.permute(dims)
|
|
|
|
|
|
def _diag(x: Tensor):
|
|
"""
|
|
like torch diag(), this returns the diagonal of a matrix; but it returns an
|
|
aliased tensor, and it supports batch dims,
|
|
i.e. input of shape (B, M, M) returns
|
|
output of shape (B, M), or input of shape (A, B, M, M) returns output of shape
|
|
(A, B, M).
|
|
"""
|
|
stride = x.stride()
|
|
if x.ndim == 3:
|
|
(B, M, M2) = x.shape
|
|
assert M == M2
|
|
ans = x.as_strided(size=(B, M), stride=(stride[0], stride[1] + stride[2]))
|
|
elif x.ndim == 4:
|
|
(B, C, M, M2) = x.shape
|
|
assert M == M2
|
|
ans = x.as_strided(size=(B, C, M), stride=(stride[0], stride[1], stride[2] + stride[3]))
|
|
elif x.ndim == 2:
|
|
(M, M2) = x.shape
|
|
assert M == M2
|
|
ans = x.as_strided(size=(M,), stride=(stride[0] + stride[1],))
|
|
return ans
|
|
|
|
|
|
def _sum(x: Tensor,
|
|
exclude_dims: List[int] = [],
|
|
keepdim: bool = False):
|
|
"""
|
|
Version of torch.sum where you specify dims to exclude, not dims to sum.
|
|
"""
|
|
if len(exclude_dims) == 0:
|
|
# dim=[] works the same as tuple(range(x.ndim)), which makes
|
|
# no sense.. supplying the full tuple for clarity and in case
|
|
# the interface changes in future.
|
|
return x.sum(dim=tuple(range(x.ndim)), keepdim=keepdim)
|
|
elif x.ndim == 1:
|
|
assert exclude_dim == [0] or exclude_dim == [-1]
|
|
return x # if one dim is excluded, there are no dims to sum, and
|
|
# x.sum(dim=[]) sums all dims so we should not call sum().
|
|
exclude_dims_norm = [i if i >= 0 else i + x.ndim for i in exclude_dims]
|
|
dims = [i for i in range(x.ndim) if i not in exclude_dims_norm]
|
|
assert (len(dims) + len(exclude_dims) == x.ndim), exclude_dims
|
|
return x.sum(dim=dims, keepdim=keepdim)
|
|
|
|
|
|
def _mean(x: Tensor,
|
|
exclude_dims: List[int] = [],
|
|
keepdim: bool = False):
|
|
"""
|
|
Version of torch.mean where you specify dims to exclude, not dims to mean.
|
|
"""
|
|
if len(exclude_dims) == 0:
|
|
# dim=[] works the same as tuple(range(x.ndim)), which makes
|
|
# no sense.. supplying the full tuple for clarity and in case
|
|
# the interface changes in future.
|
|
return x.mean(dim=tuple(range(x.ndim)), keepdim=keepdim)
|
|
elif x.ndim == 1:
|
|
assert exclude_dim == [0] or exclude_dim == [-1]
|
|
return x # if one dim is excluded, there are no dims to mean, and
|
|
# x.mean(dim=[]) means all dims so we should not call mean().
|
|
exclude_dims_norm = [i if i >= 0 else i + x.ndim for i in exclude_dims]
|
|
dims = [i for i in range(x.ndim) if i not in exclude_dims_norm]
|
|
assert (len(dims) + len(exclude_dims) == x.ndim), exclude_dims
|
|
return x.mean(dim=dims, keepdim=keepdim)
|
|
|
|
|
|
|
|
|
|
def _soft_floor(x: Tensor, floor: Union[float, Tensor], depth: float = 3) -> Tensor:
|
|
"""
|
|
Applies a floor `floor` to x in a soft way such that values of x below the floor,
|
|
as low as floor**2, stay reasonably distinct. `floor` will be mapped to
|
|
3*floor,
|
|
|
|
sqrt(floor), 0 will be mapped to `floor`, floor**2 will be mapped to 2*floor,
|
|
and values above about sqrt(floor) will be approximately linear.
|
|
|
|
Args:
|
|
x : Tensor, a Tensor of any shape, to apply a floor to, will not be modified.
|
|
Its values must be >= 0.
|
|
floor: the floor to apply, must be >0. E.g. 1.0e-10
|
|
depth: the number of iterations in the algorithm; more means it will map
|
|
tiny values more aggressively to larger values. Recommend
|
|
not more than 5.
|
|
"""
|
|
terms = [x]
|
|
cur_pow = 1.0
|
|
for i in range(depth):
|
|
x = x.sqrt()
|
|
cur_pow *= 0.5
|
|
# the reason for dividing by `depth` is that regardless of the depth,
|
|
# we want `floor` itself to be mapped to 2*floor by the algorithm,
|
|
# so we want this sum to eventually equal `floor`, if x were equal
|
|
# to `floor`
|
|
terms.append(x * ((floor ** (1-cur_pow)) / depth))
|
|
ans = torch.stack(terms)
|
|
ans = ans.sum(dim=0)
|
|
return ans + floor
|
|
|
|
|
|
|
|
def _svd(x: Tensor):
|
|
# Wrapper for torch svd that catches errors and re-tries (to address bugs in
|
|
# some versions of torch SVD for CUDA)
|
|
randU = None
|
|
for i in range(4):
|
|
try:
|
|
U, S, V = x.svd()
|
|
s = U.sum()
|
|
if s - s != 0 or (__name__ == '__main__' and random.random() < 0.01):
|
|
raise RuntimeError(f"svd failed (or random test), sum={s}")
|
|
if randU is not None:
|
|
U = torch.matmul(randU.t(), U)
|
|
return U, S, V # success
|
|
except:
|
|
logging.warning(f"svd failed: i={i}, x.shape={tuple(x.shape)}, x.sum()={x.sum()}: likely "
|
|
f"error in SVD. Will try again after random change.")
|
|
this_x = x
|
|
while this_x.ndim > 2:
|
|
this_x = this_x[0] # get rid of any batch dims
|
|
U, S, V = torch.randn_like(this_x).svd()
|
|
x = torch.matmul(U, x)
|
|
if randU is None:
|
|
randU = U
|
|
else:
|
|
randU = torch.matmul(U, randU)
|
|
|
|
|
|
|
|
class Cain(Optimizer):
|
|
"""Implements Cain algorithm, which is a substantially modified form of the
|
|
Adam optimizer.
|
|
The modifications can be summarized as follows:
|
|
|
|
(i) change Adam to what we called "Eve", which adds an extra configuration
|
|
value, "target_rms" (default value: 0.1); and for non-scalar parameters,
|
|
if their root-mean-square value exceeds "target_rms", we apply weight
|
|
decay (aggressive enough weight decay that the rms stays at target_rms).
|
|
This weight decay is applied in the same way as in AdamW, i.e. by
|
|
parameter shrinkage rather than by modifying the gradients.
|
|
(ii) By limiting the rms to 0.1, we reduce our modeling power; we restore
|
|
it by replacing all weights and biases-- in fact, all model parameters--
|
|
with expressions like
|
|
weight * weight_scale.exp() or bias * bias_scale.exp().
|
|
This has the benefit of making the learnable offsets and scales of
|
|
BatchNorm/LayerNorm unnecessary.
|
|
Below we describe the changes from Eve to Cain:
|
|
(iii) We removed all the weight_scale and bias_scale parameters from
|
|
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.
|
|
(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).
|
|
|
|
The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
|
|
The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
|
|
Eve is unpublished so far.
|
|
|
|
Arguments:
|
|
params (iterable): iterable of parameters to optimize or dicts defining
|
|
parameter groups
|
|
lr (float, optional): learning rate (default: 1e-3)
|
|
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)
|
|
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
|
|
will not go much below this.
|
|
param_max (float, optional): maximum absolute value for any parameter;
|
|
we will clamp parameters to satisfy -param_max <= param <= param_max
|
|
|
|
|
|
|
|
.. _Adam\: A Method for Stochastic Optimization:
|
|
https://arxiv.org/abs/1412.6980
|
|
.. _Decoupled Weight Decay Regularization:
|
|
https://arxiv.org/abs/1711.05101
|
|
.. _On the Convergence of Adam and Beyond:
|
|
https://openreview.net/forum?id=ryQu7f-RZ
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
params,
|
|
lr=1e-2,
|
|
betas=(0.9, 0.98),
|
|
eps=1e-8,
|
|
scale_speed=0.05,
|
|
rms_eps=1.0e-05,
|
|
param_max=20.0,
|
|
):
|
|
if not 0.0 <= lr:
|
|
raise ValueError("Invalid learning rate: {}".format(lr))
|
|
if not 0.0 <= eps:
|
|
raise ValueError("Invalid epsilon value: {}".format(eps))
|
|
if not 0.0 <= betas[0] < 1.0:
|
|
raise ValueError(
|
|
"Invalid beta parameter at index 0: {}".format(betas[0])
|
|
)
|
|
if not 0.0 <= betas[1] < 1.0:
|
|
raise ValueError(
|
|
"Invalid beta parameter at index 1: {}".format(betas[1])
|
|
)
|
|
if not 0 < scale_speed < 1.0:
|
|
raise ValueError("Invalid scale_speed value: {}".format(scale_speed))
|
|
if not 0.1 < param_max <= 10000.0:
|
|
raise ValueError("Invalid param_max value: {}".format(param_max))
|
|
if not 0.0 < rms_eps < 0.1:
|
|
raise ValueError("Invalid rms_eps value: {}".format(rms_eps))
|
|
|
|
defaults = dict(
|
|
lr=lr,
|
|
betas=betas,
|
|
eps=eps,
|
|
scale_speed=0.05,
|
|
rms_eps=rms_eps,
|
|
param_max=param_max,
|
|
)
|
|
super(Cain, self).__init__(params, defaults)
|
|
|
|
def __setstate__(self, state):
|
|
super(Cain, self).__setstate__(state)
|
|
|
|
@torch.no_grad()
|
|
def step(self, closure=None):
|
|
"""Performs a single optimization step.
|
|
|
|
Arguments:
|
|
closure (callable, optional): A closure that reevaluates the model
|
|
and returns the loss.
|
|
"""
|
|
loss = None
|
|
if closure is not None:
|
|
with torch.enable_grad():
|
|
loss = closure()
|
|
|
|
for group in self.param_groups:
|
|
beta1, beta2 = group["betas"]
|
|
lr = group["lr"]
|
|
scale_speed = group["scale_speed"]
|
|
rms_eps = group["rms_eps"]
|
|
eps = group["eps"]
|
|
param_max = group["param_max"]
|
|
|
|
for p in group["params"]:
|
|
if p.grad is None:
|
|
continue
|
|
|
|
# Perform optimization step
|
|
grad = p.grad
|
|
if grad.is_sparse:
|
|
raise RuntimeError(
|
|
"Cain does not support sparse gradients"
|
|
)
|
|
|
|
state = self.state[p]
|
|
|
|
# State initialization
|
|
if len(state) == 0:
|
|
state["step"] = 0
|
|
# Exponential moving average of gradient values
|
|
state["delta"] = torch.zeros_like(
|
|
p, memory_format=torch.preserve_format
|
|
)
|
|
# Exponential moving average of squared gradient values
|
|
state["exp_avg_sq"] = torch.zeros_like(
|
|
p, memory_format=torch.preserve_format
|
|
)
|
|
if p.numel() > 1:
|
|
# we keep track of the RMS value of the parameters to
|
|
# help set the learning rate... this emulates Eve.
|
|
state["param_rms"] = torch.ones_like(p)
|
|
# we learn the scale on the parameters in a way that
|
|
# also emulates Eve. scale_exp_avg_sq is the
|
|
# moving-average square of the gradient w.r.t. this
|
|
# scalar scale.
|
|
state["scale_exp_avg_sq"] = torch.zeros((), device=p.device,
|
|
dtype=p.dtype)
|
|
|
|
|
|
delta, exp_avg_sq = state["delta"], state["exp_avg_sq"]
|
|
step = state["step"]
|
|
|
|
delta.mul_(beta1)
|
|
|
|
if p.numel() > 1:
|
|
|
|
if True:
|
|
# This block learns the scale on the parameters.
|
|
# scale_deriv is the derivative w.r.t. a log-space scale
|
|
# on the parameters, i.e. if we imagine: p =
|
|
# underlying_param * scale.exp(), scale_deriv is the
|
|
# derivative w.r.t. scale (it does not matter
|
|
# what value `scale` currently has).
|
|
scale_deriv = (p * grad).sum()
|
|
scale_exp_avg_sq = state["scale_exp_avg_sq"]
|
|
scale_exp_avg_sq.mul_(beta2).addcmul_(scale_deriv, scale_deriv,
|
|
value=1 - beta2)
|
|
|
|
# should actually be step + 1, so on 1st minibatch we are not learning
|
|
# anything here. May fix this at some point.
|
|
scale_bias_correction2 = 1 - beta2 ** (step + 1)
|
|
|
|
scale_denom = (scale_exp_avg_sq.sqrt()).add_(eps)
|
|
|
|
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(),
|
|
# delta is the change in `scale`.
|
|
scale_delta = scale_alpha * (scale_deriv / scale_denom)
|
|
|
|
delta.add_(p, alpha=scale_delta)
|
|
|
|
# Decay the second moment running average coefficient
|
|
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
|
|
|
|
bias_correction2 = 1 - beta2 ** step
|
|
|
|
grad_rms = (exp_avg_sq.sqrt()).add_(eps)
|
|
|
|
this_delta = grad / grad_rms
|
|
|
|
param_rms = state["param_rms"]
|
|
self._update_param_rms(p, param_rms, step, rms_eps)
|
|
|
|
alpha = (-(1-beta1) * lr * bias_correction2)
|
|
|
|
# geometrically interpolate the per-element param_rms with its mean
|
|
#param_rms_half = (param_rms * param_rms.mean()) ** 0.5
|
|
delta.add_(this_delta * param_rms, alpha=alpha)
|
|
|
|
# below, will do: delta.add_(this_delta, alpha=alpha)
|
|
else:
|
|
# p.numel() == 1
|
|
|
|
# Decay the second moment running average coefficient
|
|
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
|
|
bias_correction2 = 1 - beta2 ** step
|
|
denom = (exp_avg_sq.sqrt()).add_(eps)
|
|
this_delta = grad / denom
|
|
alpha = -lr * (1-beta1) * (bias_correction2 ** 0.5)
|
|
delta.add_(this_delta, alpha=alpha)
|
|
|
|
p.add_(delta)
|
|
if step % 10 == 0:
|
|
p.clamp_(min=-param_max, max=param_max)
|
|
state["step"] += 1
|
|
|
|
return loss
|
|
|
|
def _update_param_rms(self, p: Tensor, param_rms: Tensor,
|
|
step: int, rms_eps: float):
|
|
"""
|
|
Updates `param_rms`. Require p.numel() > 1
|
|
Args:
|
|
p: parameter whose magnitude/rms we are measuring
|
|
param_rms: a Tensor containing non-negative values, with the same
|
|
shape as p. Contains an estimate of the root-mean-square
|
|
value of the parameter. This has a special structure where
|
|
it must be a product of one-dimensional quantities,
|
|
one for each axis of p.
|
|
We exclude certain axes to avoid estimating the rms value
|
|
from a too-small number of parameters.
|
|
step: the step of the algorithm
|
|
rms_eps: epsilon such that we'll aim to prevent param_rms from
|
|
going much lower than this.
|
|
"""
|
|
if step % 1000 == 0:
|
|
# Periodically refresh param_rms to keep the invariance that it is a
|
|
# product over its axes (numerical roundoff might destroy this)
|
|
param_rms.fill_(1.0)
|
|
for dim in range(p.ndim):
|
|
self._update_param_rms_for_dim(p, param_rms, dim, rms_eps)
|
|
else:
|
|
dim = step % p.ndim
|
|
self._update_param_rms_for_dim(p, param_rms, dim, rms_eps)
|
|
|
|
def _update_param_rms_for_dim(self, p: Tensor,
|
|
param_rms: Tensor,
|
|
dim: int,
|
|
rms_eps: float):
|
|
"""
|
|
Updates `param_rms` on one of its axes/dimensions.
|
|
Args:
|
|
p: parameter whose magnitude/rms we are measuring
|
|
param_rms: a Tensor containing non-negative values, with the same
|
|
shape as p. Contains an estimate of the root-mean-square
|
|
value of the parameter. This has a special structure where
|
|
it must be a product of one-dimensional quantities,
|
|
one for each axis of p.
|
|
We exclude certain axes to avoid estimating the rms value
|
|
from a too-small number of parameters.
|
|
dim: with 0 <= dim < p.ndim, the
|
|
dimension/axis on which to update the scale factor
|
|
rms_eps: epsilon such that we'll aim to prevent param_rms from
|
|
going much lower than this.
|
|
"""
|
|
# var_factor is the square of the factor to change param_rms by. p will
|
|
# not be super large, we limit it to -param_max <= p <= param_max.
|
|
# Clamping p**2 to min=1.0e-10 should prevent param_rms from getting much
|
|
# smaller than 1.0e-05, which should prevent division by zero here.
|
|
var_factor = (p ** 2).clamp(min=rms_eps*rms_eps) / (param_rms ** 2)
|
|
|
|
if p.numel() <= 2 * p.shape[dim]:
|
|
var_factor = var_factor.mean()
|
|
else:
|
|
dims = tuple(d for d in range(p.ndim) if d != dim)
|
|
var_factor = var_factor.mean(dim=dims, keepdim=True)
|
|
|
|
#if random.random() < 0.01:
|
|
# logging.info(f"shape={p.shape}, var_factor={var_factor}")
|
|
param_rms.mul_(var_factor.sqrt())
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class LRScheduler(object):
|
|
"""
|
|
Base-class for learning rate schedulers where the learning-rate depends on both the
|
|
batch and the epoch.
|
|
"""
|
|
|
|
def __init__(self, optimizer: Optimizer, verbose: bool = False):
|
|
# Attach optimizer
|
|
if not isinstance(optimizer, Optimizer):
|
|
raise TypeError(
|
|
"{} is not an Optimizer".format(type(optimizer).__name__)
|
|
)
|
|
self.optimizer = optimizer
|
|
self.verbose = verbose
|
|
|
|
for group in optimizer.param_groups:
|
|
group.setdefault("initial_lr", group["lr"])
|
|
|
|
self.base_lrs = [
|
|
group["initial_lr"] for group in optimizer.param_groups
|
|
]
|
|
|
|
self.epoch = 0
|
|
self.batch = 0
|
|
|
|
def state_dict(self):
|
|
"""Returns the state of the scheduler as a :class:`dict`.
|
|
|
|
It contains an entry for every variable in self.__dict__ which
|
|
is not the optimizer.
|
|
"""
|
|
return {
|
|
"base_lrs": self.base_lrs,
|
|
"epoch": self.epoch,
|
|
"batch": self.batch,
|
|
}
|
|
|
|
def load_state_dict(self, state_dict):
|
|
"""Loads the schedulers state.
|
|
|
|
Args:
|
|
state_dict (dict): scheduler state. Should be an object returned
|
|
from a call to :meth:`state_dict`.
|
|
"""
|
|
self.__dict__.update(state_dict)
|
|
|
|
def get_last_lr(self) -> List[float]:
|
|
"""Return last computed learning rate by current scheduler. Will be a list of float."""
|
|
return self._last_lr
|
|
|
|
def get_lr(self):
|
|
# Compute list of learning rates from self.epoch and self.batch and
|
|
# self.base_lrs; this must be overloaded by the user.
|
|
# e.g. return [some_formula(self.batch, self.epoch, base_lr) for base_lr in self.base_lrs ]
|
|
raise NotImplementedError
|
|
|
|
def step_batch(self, batch: Optional[int] = None) -> None:
|
|
# Step the batch index, or just set it. If `batch` is specified, it
|
|
# must be the batch index from the start of training, i.e. summed over
|
|
# all epochs.
|
|
# You can call this in any order; if you don't provide 'batch', it should
|
|
# of course be called once per batch.
|
|
if batch is not None:
|
|
self.batch = batch
|
|
else:
|
|
self.batch = self.batch + 1
|
|
self._set_lrs()
|
|
|
|
def step_epoch(self, epoch: Optional[int] = None):
|
|
# Step the epoch index, or just set it. If you provide the 'epoch' arg,
|
|
# you should call this at the start of the epoch; if you don't provide the 'epoch'
|
|
# arg, you should call it at the end of the epoch.
|
|
if epoch is not None:
|
|
self.epoch = epoch
|
|
else:
|
|
self.epoch = self.epoch + 1
|
|
self._set_lrs()
|
|
|
|
def _set_lrs(self):
|
|
values = self.get_lr()
|
|
assert len(values) == len(self.optimizer.param_groups)
|
|
|
|
for i, data in enumerate(zip(self.optimizer.param_groups, values)):
|
|
param_group, lr = data
|
|
param_group["lr"] = lr
|
|
self.print_lr(self.verbose, i, lr)
|
|
self._last_lr = [group["lr"] for group in self.optimizer.param_groups]
|
|
|
|
def print_lr(self, is_verbose, group, lr):
|
|
"""Display the current learning rate."""
|
|
if is_verbose:
|
|
logging.info(
|
|
f"Epoch={self.epoch}, batch={self.batch}: adjusting learning rate"
|
|
f" of group {group} to {lr:.4e}."
|
|
)
|
|
|
|
class Eve(Optimizer):
|
|
"""
|
|
Implements Eve algorithm. This is a modified version of AdamW with a special
|
|
way of setting the weight-decay / shrinkage-factor, which is designed to make the
|
|
rms of the parameters approach a particular target_rms (default: 0.1). This is
|
|
for use with networks with 'scaled' versions of modules (see scaling.py), which
|
|
will be close to invariant to the absolute scale on the parameter matrix.
|
|
|
|
The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
|
|
The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
|
|
Eve is unpublished so far.
|
|
|
|
Arguments:
|
|
params (iterable): iterable of parameters to optimize or dicts defining
|
|
parameter groups
|
|
lr (float, optional): learning rate (default: 1e-3)
|
|
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)
|
|
weight_decay (float, optional): weight decay coefficient (default: 3e-4;
|
|
this value means that the weight would decay significantly after
|
|
about 3k minibatches. Is not multiplied by learning rate, but
|
|
is conditional on RMS-value of parameter being > target_rms.
|
|
target_rms (float, optional): target root-mean-square value of
|
|
parameters, if they fall below this we will stop applying weight decay.
|
|
|
|
|
|
.. _Adam\: A Method for Stochastic Optimization:
|
|
https://arxiv.org/abs/1412.6980
|
|
.. _Decoupled Weight Decay Regularization:
|
|
https://arxiv.org/abs/1711.05101
|
|
.. _On the Convergence of Adam and Beyond:
|
|
https://openreview.net/forum?id=ryQu7f-RZ
|
|
"""
|
|
def __init__(
|
|
self,
|
|
params,
|
|
lr=1e-3,
|
|
betas=(0.9, 0.98),
|
|
eps=1e-8,
|
|
weight_decay=1e-3,
|
|
target_rms=0.1,
|
|
):
|
|
if not 0.0 <= lr:
|
|
raise ValueError("Invalid learning rate: {}".format(lr))
|
|
if not 0.0 <= eps:
|
|
raise ValueError("Invalid epsilon value: {}".format(eps))
|
|
if not 0.0 <= betas[0] < 1.0:
|
|
raise ValueError(
|
|
"Invalid beta parameter at index 0: {}".format(betas[0])
|
|
)
|
|
if not 0.0 <= betas[1] < 1.0:
|
|
raise ValueError(
|
|
"Invalid beta parameter at index 1: {}".format(betas[1])
|
|
)
|
|
if not 0 <= weight_decay <= 0.1:
|
|
raise ValueError(
|
|
"Invalid weight_decay value: {}".format(weight_decay)
|
|
)
|
|
if not 0 < target_rms <= 10.0:
|
|
raise ValueError("Invalid target_rms value: {}".format(target_rms))
|
|
defaults = dict(
|
|
lr=lr,
|
|
betas=betas,
|
|
eps=eps,
|
|
weight_decay=weight_decay,
|
|
target_rms=target_rms,
|
|
)
|
|
super(Eve, self).__init__(params, defaults)
|
|
|
|
def __setstate__(self, state):
|
|
super(Eve, self).__setstate__(state)
|
|
|
|
@torch.no_grad()
|
|
def step(self, closure=None):
|
|
"""Performs a single optimization step.
|
|
|
|
Arguments:
|
|
closure (callable, optional): A closure that reevaluates the model
|
|
and returns the loss.
|
|
"""
|
|
loss = None
|
|
if closure is not None:
|
|
with torch.enable_grad():
|
|
loss = closure()
|
|
|
|
for group in self.param_groups:
|
|
for p in group["params"]:
|
|
if p.grad is None:
|
|
continue
|
|
|
|
# Perform optimization step
|
|
grad = p.grad
|
|
if grad.is_sparse:
|
|
raise RuntimeError(
|
|
"AdamW does not support sparse gradients"
|
|
)
|
|
|
|
state = self.state[p]
|
|
|
|
# State initialization
|
|
if len(state) == 0:
|
|
state["step"] = 0
|
|
# Exponential moving average of gradient values
|
|
state["exp_avg"] = torch.zeros_like(
|
|
p, memory_format=torch.preserve_format
|
|
)
|
|
# Exponential moving average of squared gradient values
|
|
state["exp_avg_sq"] = torch.zeros_like(
|
|
p, memory_format=torch.preserve_format
|
|
)
|
|
|
|
exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
|
|
|
|
beta1, beta2 = group["betas"]
|
|
|
|
state["step"] += 1
|
|
bias_correction1 = 1 - beta1 ** state["step"]
|
|
bias_correction2 = 1 - beta2 ** state["step"]
|
|
|
|
# Decay the first and second moment running average coefficient
|
|
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
|
|
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
|
|
denom = (exp_avg_sq.sqrt() * (bias_correction2 ** -0.5)).add_(
|
|
group["eps"]
|
|
)
|
|
|
|
step_size = group["lr"] / bias_correction1
|
|
target_rms = group["target_rms"]
|
|
weight_decay = group["weight_decay"]
|
|
|
|
if p.numel() > 1:
|
|
# avoid applying this weight-decay on "scaling factors"
|
|
# (which are scalar).
|
|
is_above_target_rms = p.norm() > (
|
|
target_rms * (p.numel() ** 0.5)
|
|
)
|
|
p.mul_(1 - (weight_decay * is_above_target_rms))
|
|
|
|
p.addcdiv_(exp_avg, denom, value=-step_size)
|
|
|
|
if random.random() < 0.0005:
|
|
step = (exp_avg/denom) * step_size
|
|
logging.info(f"Delta rms = {(step**2).mean().item()}, shape = {step.shape}")
|
|
|
|
|
|
return loss
|
|
|
|
|
|
class Eden(LRScheduler):
|
|
"""
|
|
Eden scheduler.
|
|
lr = initial_lr * (((batch**2 + lr_batches**2) / lr_batches**2) ** -0.25 *
|
|
(((epoch**2 + lr_epochs**2) / lr_epochs**2) ** -0.25))
|
|
|
|
E.g. suggest initial-lr = 0.003 (passed to optimizer).
|
|
|
|
Args:
|
|
optimizer: the optimizer to change the learning rates on
|
|
lr_batches: the number of batches after which we start significantly
|
|
decreasing the learning rate, suggest 5000.
|
|
lr_epochs: the number of epochs after which we start significantly
|
|
decreasing the learning rate, suggest 6 if you plan to do e.g.
|
|
20 to 40 epochs, but may need smaller number if dataset is huge
|
|
and you will do few epochs.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
lr_batches: Union[int, float],
|
|
lr_epochs: Union[int, float],
|
|
verbose: bool = False,
|
|
):
|
|
super(Eden, self).__init__(optimizer, verbose)
|
|
self.lr_batches = lr_batches
|
|
self.lr_epochs = lr_epochs
|
|
|
|
def get_lr(self):
|
|
factor = (
|
|
(self.batch ** 2 + self.lr_batches ** 2) / self.lr_batches ** 2
|
|
) ** -0.25 * (
|
|
((self.epoch ** 2 + self.lr_epochs ** 2) / self.lr_epochs ** 2)
|
|
** -0.25
|
|
)
|
|
return [x * factor for x in self.base_lrs]
|
|
|
|
|
|
def _test_eden():
|
|
m = torch.nn.Linear(100, 100)
|
|
optim = Cain(m.parameters(), lr=0.003)
|
|
|
|
scheduler = Eden(optim, lr_batches=100, lr_epochs=2, verbose=True)
|
|
|
|
for epoch in range(10):
|
|
scheduler.step_epoch(epoch) # sets epoch to `epoch`
|
|
|
|
for step in range(20):
|
|
x = torch.randn(200, 100).detach()
|
|
x.requires_grad = True
|
|
y = m(x)
|
|
dy = torch.randn(200, 100).detach()
|
|
f = (y * dy).sum()
|
|
f.backward()
|
|
|
|
optim.step()
|
|
scheduler.step_batch()
|
|
optim.zero_grad()
|
|
|
|
logging.info(f"last lr = {scheduler.get_last_lr()}")
|
|
logging.info(f"state dict = {scheduler.state_dict()}")
|
|
|
|
|
|
def _test_eve_cain():
|
|
import timeit
|
|
from scaling import ScaledLinear
|
|
E = 100
|
|
B = 4
|
|
T = 2
|
|
logging.info("in test_eve_cain")
|
|
#device = torch.device('cuda')
|
|
device = torch.device('cpu')
|
|
dtype = torch.float32
|
|
|
|
fix_random_seed(42)
|
|
# these input_magnitudes and output_magnitudes are to test that
|
|
# Abel is working as we expect and is able to adjust scales of
|
|
# different dims differently.
|
|
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 [3, 2, 1, 0]: # will restore 1,0 later
|
|
for iter in [3, 2]:
|
|
fix_random_seed(42)
|
|
Linear = torch.nn.Linear if iter == 0 else ScaledLinear
|
|
|
|
hidden_dim = 768
|
|
m = torch.nn.Sequential(Linear(E, hidden_dim),
|
|
torch.nn.PReLU(),
|
|
Linear(hidden_dim, hidden_dim),
|
|
torch.nn.PReLU(),
|
|
Linear(hidden_dim, E),
|
|
).to(device)
|
|
|
|
train_pairs = [ (100.0 * torch.randn(B, T, E, device=device, dtype=dtype) * input_magnitudes,
|
|
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.03)
|
|
elif iter == 2: optim = PrAdam(m.parameters(), lr=0.03)
|
|
elif iter == 3: optim = PrAdam(m.parameters(), lr=0.03, max_block_size=256)
|
|
scheduler = Eden(optim, lr_batches=200, lr_epochs=5, verbose=False)
|
|
|
|
#TEMP
|
|
if iter == 3:
|
|
a = torch.randn(5, 10, 5, 5)
|
|
a = torch.matmul(a, a.transpose(2, 3))
|
|
b1 = optim._smooth_cov(a, 0.1, 4.0, 0.999999999)
|
|
b2 = optim._smooth_cov(a, 0.1, 4.0, 1.0)
|
|
diff = (b1 - b2)
|
|
ratio = (diff**2).sqrt().mean() / (b1**2).sqrt().mean()
|
|
logging.info(f"ratio = {ratio}")
|
|
assert ratio < 0.01
|
|
|
|
start = timeit.default_timer()
|
|
avg_loss = 0.0
|
|
for epoch in range(180):
|
|
scheduler.step_epoch()
|
|
#if epoch == 100 and iter in [2,3]:
|
|
# optim.reset_speedup() # check it doesn't crash.
|
|
|
|
#if epoch == 130:
|
|
# opts = diagnostics.TensorDiagnosticOptions(
|
|
# 2 ** 22
|
|
# ) # allow 4 megabytes per sub-module
|
|
# diagnostic = diagnostics.attach_diagnostics(m, opts)
|
|
|
|
|
|
for n, (x,y) in enumerate(train_pairs):
|
|
y_out = m(x)
|
|
loss = ((y_out - y)**2).mean() * 100.0
|
|
if epoch == 0 and n == 0:
|
|
avg_loss = loss.item()
|
|
else:
|
|
avg_loss = 0.98 * avg_loss + 0.02 * loss.item()
|
|
if n == 0 and epoch % 5 == 0:
|
|
#norm1 = '%.2e' % (m[0].weight**2).mean().sqrt().item()
|
|
#norm1b = '%.2e' % (m[0].bias**2).mean().sqrt().item()
|
|
#norm2 = '%.2e' % (m[2].weight**2).mean().sqrt().item()
|
|
#norm2b = '%.2e' % (m[2].bias**2).mean().sqrt().item()
|
|
#scale1 = '%.2e' % (m[0].weight_scale.exp().item())
|
|
#scale1b = '%.2e' % (m[0].bias_scale.exp().item())
|
|
#scale2 = '%.2e' % (m[2].weight_scale.exp().item())
|
|
#scale2b = '%.2e' % (m[2].bias_scale.exp().item())
|
|
lr = scheduler.get_last_lr()[0]
|
|
logging.info(f"Iter {iter}, epoch {epoch}, batch {n}, avg_loss {avg_loss:.4g}, lr={lr:.4e}") #, norms={norm1,norm1b,norm2,norm2b}") # scales={scale1,scale1b,scale2,scale2b}
|
|
loss.log().backward()
|
|
optim.step()
|
|
optim.zero_grad()
|
|
scheduler.step_batch()
|
|
|
|
#diagnostic.print_diagnostics()
|
|
|
|
stop = timeit.default_timer()
|
|
logging.info(f"Iter={iter}, Time taken: {stop - start}")
|
|
|
|
logging.info(f"last lr = {scheduler.get_last_lr()}")
|
|
#logging.info("state dict = ", scheduler.state_dict())
|
|
#logging.info("optim state_dict = ", optim.state_dict())
|
|
logging.info(f"input_magnitudes = {input_magnitudes}")
|
|
logging.info(f"output_magnitudes = {output_magnitudes}")
|
|
|
|
|
|
def _test_svd():
|
|
device = 'cuda'
|
|
for i in range(1000):
|
|
X = torch.eye(100, device=device) + 1.0e-08 * torch.randn(100, 100, device=device)
|
|
U, S, V = _svd(X)
|
|
X2 = torch.matmul(U*S, V.t())
|
|
assert torch.allclose(X2, X, atol=0.001)
|
|
|
|
def _test_soft_floor():
|
|
x = torch.arange(100) - 50
|
|
y = x.to(torch.float32).exp()
|
|
z = _soft_floor(y, 1.0e-06)
|
|
print(f"y={y}, z={z}")
|
|
|
|
z = _soft_floor(torch.tensor(1.0e-06), 1.0e-06)
|
|
print("z = ", z)
|
|
assert (z - 3.0e-06).abs() < 1.0e-07
|
|
|
|
if __name__ == "__main__":
|
|
torch.set_num_threads(1)
|
|
torch.set_num_interop_threads(1)
|
|
logging.getLogger().setLevel(logging.INFO)
|
|
import subprocess
|
|
s = subprocess.check_output("git status -uno .; git log -1", shell=True)
|
|
logging.info(s)
|
|
#_test_svd()
|
|
_test_soft_floor()
|
|
_test_eve_cain()
|
|
#_test_eden()
|