# coding=utf-8

# Copyright 2018 The Tensor2Tensor 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.

"""Optimization."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

# Dependency imports

import tensorflow as tf

class AdafactorOptimizer(tf.train.Optimizer):

"""Optimizer that implements the Adafactor algorithm.

Adafactor is described in https://arxiv.org/abs/1804.04235.

Adafactor is most similar to Adam (Kingma and Ba), the major differences are:

1. For a two-dimensional AxB weight matrix, Adafactor uses only A+B auxiliary

parameters to maintain the second-moment estimator, instead of AB.

This is advantageous on memory-limited systems. In addition, beta1

(momentum) is set to zero by default, saving an additional auxiliary

parameter per weight.

2. Adafactor incorporates "update-clipping" - a scale-invariant analog of

gradient clipping. This adds stability

3. Adafactor does not require an external "learning rate". By default, it

incorporates a relative-update-scale schedule, corresponding to

inverse-square-root learning-rate-decay in ADAM. We hope this works well

for most applications.

ALGORITHM:

parameter -= absolute_update_scale * clip(grad / grad_scale)

where:

absolute_update_scale := relative_update_scale * parameter_scale

relative_update_scale := min((step_num + 1)**-0.5, 1e-2)

parameter_scale := max(rms(var)), 1e-3)

clip(x) := x / max(1.0, rms(x))

grad_scale := tf.sqrt(v) (v is the second-moment estimator)

The second-moment estimator v is maintained in a manner similar to Adam:

We initialize

```

if var is 2-dimensional:

v_r <- zeros([num_rows])

v_c <- zeros([num_cols])

else:

v <- zeros(shape(var))

```

The update rule is as follows:

```

decay_rate = 1 - (step_num + 1) ^ -0.8

grad_squared = tf.square(grad) + epsilon

if var is 2-dimensional:

v_r <- decay_rate * v_r + (1 - decay_rate) * reduce_mean(grad_squared, 1)

v_c <- decay_rate * v_c + (1 - decay_rate) * reduce_mean(grad_squared, 0)

v = outer_prod(v_r, v_c) / reduce_mean(v_r)

else:

v <- decay_rate * v + (1 - decay_rate) * grad_squared

```

Several parts of this algorithm are configurable from the initializer.

multiply_by_parameter_scale: If True, then compute absolute_update_scale

as described above. If False, let absolute_update_scale be the externally

supplied learning_rate.

learning_rate: represents relative_update_scale if

multiply_by_parameter_scale==True, or absolute_update_scale if

multiply_by_parameter_scale==False.

decay_rate: Decay rate of the second moment estimator (varies by step_num).

This should be set to a function such that:

1-1/(step_num + 1) <= decay_rate(step_num) < 1.0

beta1: enables momentum, as in Adam. Uses extra memory if nonzero.

clipping_threshold: should be >=1.0 or None for no update clipping

factored: whether to factor the second-moment estimator. True means

less memory usage.

TODO(noam): we should also apply the 2d logic to the two final dimensions.

of >2d convolutional kernels.

"""

def __init__(self,

multiply_by_parameter_scale=True,

learning_rate=None,

decay_rate=None,

beta1=0.0,

clipping_threshold=1.0,

factored=True,

simulated_quantize_bits=None,

use_locking=False,

name="Adafactor"):

"""Construct a new Adafactor optimizer.

See class comment.

Args:

multiply_by_parameter_scale: a boolean

learning_rate: an optional Scalar.

decay_rate: an optional Scalar.

beta1: a float value between 0 and 1

clipping_threshold: an optional float >= 1

factored: a boolean - whether to use factored second-moment estimator

for 2d variables

simulated_quantize_bits: train with simulated quantized parameters

(experimental)

use_locking: If True use locks for update operations.

name: Optional name for the operations created when applying gradients.

Defaults to "AdafactorOptimizer".

Raises:

ValueError: if absolute_update_scale and relative_update_scale_fn are both

present or both absent.

"""

super(AdafactorOptimizer, self).__init__(use_locking, name)

self._multiply_by_parameter_scale = multiply_by_parameter_scale

if learning_rate is None:

learning_rate = self._learning_rate_default(multiply_by_parameter_scale)

self._learning_rate = learning_rate

if decay_rate is None:

decay_rate = self._decay_rate_default()

self._decay_rate = decay_rate

self._beta1 = beta1

self._clipping_threshold = clipping_threshold

self._factored = factored

self._simulated_quantize_bits = simulated_quantize_bits

if self._simulated_quantize_bits:

self._quantization_noise = _quantization_noise_from_step_num()

def _should_use_factored_second_moment_estimate(self, shape):

"""Should we use a factored second moment estimator.

Based on the shape of the variable.

Args:

shape: a list of integers

Returns:

a boolean

"""

return self._factored and len(shape) == 2

def _create_slots(self, var_list):

for var in var_list:

shape = var.get_shape().as_list()

if self._beta1:

self._zeros_slot(var, "m", self._name)

if self._should_use_factored_second_moment_estimate(shape):

r_val = tf.zeros([shape[0]], dtype=tf.float32)

c_val = tf.zeros([shape[1]], dtype=tf.float32)

self._get_or_make_slot(var, r_val, "vr", self._name)

self._get_or_make_slot(var, c_val, "vc", self._name)

else:

v_val = tf.zeros(shape, dtype=tf.float32)

self._get_or_make_slot(var, v_val, "v", self._name)

def _apply_dense(self, grad, var):

return self._resource_apply_dense(grad, var)

def _apply_sparse(self, grad, var):

return self._apply_dense(tf.convert_to_tensor(grad), var)

def _parameter_scale(self, var):

"""Estimate the scale of the parameters from the current values.

We include a minimum value of 0.001 to give it a chance to escape 0

if it was zero-initialized.

Instead of using the value, we could impute the scale from the shape,

as initializers do.

Args:

var: a variable or Tensor.

Returns:

a Scalar

"""

return tf.maximum(reduce_rms(var), 0.001)

def _resource_apply_dense(self, grad, var):

grad = tf.to_float(grad)

grad_squared = tf.square(grad) + 1e-30

grad_squared_mean = tf.reduce_mean(grad_squared)

decay_rate = self._decay_rate

update_scale = self._learning_rate

if self._multiply_by_parameter_scale:

update_scale *= tf.to_float(self._parameter_scale(var))

# HACK: Make things dependent on grad.

# This confounds the XLA rewriter and keeps it from fusing computations

# across different variables. This fusion is a bad for HBM usage, since

# it causes the gradients to persist in memory.

decay_rate += grad_squared_mean * 1e-30

update_scale += grad_squared_mean * 1e-30

# END HACK

mixing_rate = 1.0 - decay_rate

shape = var.get_shape().as_list()

updates = []

if self._should_use_factored_second_moment_estimate(shape):

grad_squared_row_mean = tf.reduce_mean(grad_squared, 1)

grad_squared_col_mean = tf.reduce_mean(grad_squared, 0)

vr = self.get_slot(var, "vr")

new_vr = (decay_rate * vr + mixing_rate * grad_squared_row_mean)

vc = self.get_slot(var, "vc")

new_vc = (decay_rate * vc + mixing_rate * grad_squared_col_mean)

vr_update = tf.assign(vr, new_vr, use_locking=self._use_locking)

vc_update = tf.assign(vc, new_vc, use_locking=self._use_locking)

updates = [vr_update, vc_update]

long_term_mean = tf.reduce_mean(new_vr)

r_factor = tf.rsqrt(new_vr / long_term_mean)

c_factor = tf.rsqrt(new_vc)

x = grad * tf.expand_dims(r_factor, 1) * tf.expand_dims(c_factor, 0)

else:

v = self.get_slot(var, "v")

new_v = decay_rate * v + mixing_rate * grad_squared

v_update = tf.assign(v, new_v, use_locking=self._use_locking)

updates = [v_update]

x = grad * tf.rsqrt(new_v)

if self._clipping_threshold is not None:

clipping_denom = tf.maximum(1.0, reduce_rms(x) / self._clipping_threshold)

x /= clipping_denom

subtrahend = update_scale * x

if self._beta1:

m = self.get_slot(var, "m")

new_m = self._beta1 * tf.to_float(m) + (1.0 - self._beta1) * subtrahend

subtrahend = new_m

new_m = tf.cast(new_m, var.dtype)

updates.append(tf.assign(m, new_m, use_locking=self._use_locking))

new_val = tf.to_float(var) - subtrahend

if var.dtype == tf.bfloat16:

new_val = _to_bfloat16_unbiased(new_val)

if self._simulated_quantize_bits:

new_val = _simulated_quantize(

var - subtrahend, self._simulated_quantize_bits,

self._quantization_noise)

var_update = tf.assign(var, new_val, use_locking=self._use_locking)

updates = [var_update] + updates

return tf.group(*updates)

def _decay_rate_default(self):

return adafactor_decay_rate_pow(0.8)

def _learning_rate_default(self, multiply_by_parameter_scale):

learning_rate = tf.minimum(tf.rsqrt(step_num() + 1.0), 0.01)

if not multiply_by_parameter_scale:

learning_rate *= 0.05

return learning_rate

def adafactor_decay_rate_adam(beta2):

"""Second-moment decay rate like Adam, subsuming the correction factor.

Args:

beta2: a float between 0 and 1

Returns:

a scalar

"""

t = tf.to_float(tf.train.get_or_create_global_step()) + 1.0

decay = beta2 * (1.0 - tf.pow(beta2, t - 1.0)) / (1.0 - tf.pow(beta2, t))

# decay = tf.cond(tf.equal(t, 1.0), lambda: beta2, lambda: decay)

return decay

def adafactor_decay_rate_pow(exponent):

"""Second moment decay rate where memory-length grows as step_num^exponent.

Args:

exponent: a float between 0 and 1

Returns:

a scalar

"""

return 1.0 - tf.pow((step_num() + 1.0), -exponent)

def step_num():

return tf.to_float(tf.train.get_or_create_global_step())

def adafactor_optimizer_from_hparams(hparams, lr):

"""Create an Adafactor optimizer based on model hparams.

Args:

hparams: model hyperparameters

lr: learning rate scalar.

Returns:

an AdafactorOptimizer

Raises:

ValueError: on illegal values

"""

if hparams.optimizer_adafactor_decay_type == "Adam":

decay_rate = adafactor_decay_rate_adam(

hparams.optimizer_adafactor_beta2)

elif hparams.optimizer_adafactor_decay_type == "pow":

decay_rate = adafactor_decay_rate_pow(

hparams.optimizer_adafactor_memory_exponent)

else:

raise ValueError("unknown optimizer_adafactor_decay_type")

return AdafactorOptimizer(

multiply_by_parameter_scale=(

hparams.optimizer_adafactor_multiply_by_parameter_scale),

learning_rate=lr,

decay_rate=decay_rate,

beta1=hparams.optimizer_adafactor_beta1,

clipping_threshold=hparams.optimizer_adafactor_clipping_threshold,

factored=hparams.optimizer_adafactor_factored,

simulated_quantize_bits=getattr(

hparams, "simulated_parameter_quantize_bits", 0),

use_locking=False,

name="Adafactor")

def reduce_rms(x):

return tf.sqrt(tf.reduce_mean(tf.square(x)))

def _simulated_quantize(x, num_bits, quantization_noise):

"""Simulate quantization to num_bits bits, with externally-stored scale.

num_bits is the number of bits used to store each value.

quantization_noise is a float32 Tensor containing values in [0, 1).

Each value in quantization_noise should take different values across

different steps, approximating a uniform distribution over [0, 1).

In the case of replicated TPU training, quantization_noise should be identical

across replicas in order to keep the parameters identical across replicas.

The natural choice for quantization_noise would be tf.random_uniform(),

but this is not possible for TPU, since there is currently no way to seed

the different cores to produce identical values across replicas. Instead we

use _quantization_noise_from_step_num() (see below).

The quantization scheme is as follows:

Compute the maximum absolute value by row (call this max_abs).

Store this either in an auxiliary variable or in an extra column.

Divide the parameters by (max_abs / (2^(num_bits-1)-1)). This gives a

float32 value in the range [-2^(num_bits-1)-1, 2^(num_bits-1)-1]

Unbiased randomized roundoff by adding quantization_noise and rounding down.

This produces a signed integer with num_bits bits which can then be stored.

Args:

x: a float32 Tensor

num_bits: an integer between 1 and 22

quantization_noise: a float Tensor broadcastable to the shape of x.

Returns:

a float32 Tensor

"""

shape = x.get_shape().as_list()

if not (len(shape) >= 2 and shape[-1] > 1):

return x

max_abs = tf.reduce_max(tf.abs(x), -1, keep_dims=True) + 1e-9

max_int = 2 ** (num_bits - 1) - 1

scale = max_abs / max_int

x /= scale

x = tf.floor(x + quantization_noise)

# dequantize before storing (since this is a simulation)

x *= scale

return x

def _quantization_noise_from_step_num():

"""A quantization noise equal to (phi * (step_num + 1)) mod 1.0.

See _simulated_quantize.

Returns:

a float32 scalar

"""

step = tf.to_int32(tf.train.get_or_create_global_step()) + 1

phi = ((5 ** 0.5) - 1) / 2

# Naive computation tf.mod(phi * step, 1.0) in float32 would be disastrous

# due to loss of precision when the step number gets large.

# Computation in doubles does not work on TPU, so we use this complicated

# alternative computation which does not suffer from these roundoff errors.

ret = 0.0

for i in xrange(30):

ret += (((phi * (2 ** i)) % 1.0) # double-precision computation in python

* tf.to_float(tf.mod(step // (2 ** i), 2)))

return tf.mod(ret, 1.0)

def _randomized_roundoff_to_bfloat16(x, quantization_noise, cand1, cand2):

"""Round-off x to cand1 or to cand2 in an unbiased way.

Cand1 and cand2 are the same shape as x.

For every element of x, the corresponding elements of cand1 and cand2 should

be the two closest bfloat16 values to x. Order does not matter.

cand1 and cand2 must differ from each other.

Args:

x: A float32 Tensor.

quantization_noise: A Tensor broadcastable to the shape of x containing

random uniform values in [0.0, 1.0].

cand1: A bfloat16 Tensor the same shape as x.

cand2: A bfloat16 Tensor the same shape as x.

Returns:

A bfloat16 Tensor.

"""

cand1_f = tf.to_float(cand1)

cand2_f = tf.to_float(cand2)

step_size = cand2_f - cand1_f

fpart = (x - cand1_f) / step_size

ret = tf.where(tf.greater(fpart, quantization_noise), cand2, cand1)

return ret

def _to_bfloat16_unbiased(x):

"""Convert a float32 to a bfloat16 using randomized roundoff.

Note: If this ever produces worse results than using float32 all the way

through, we should try to diagnose and fix it. There are several things

to try:

1. Encode parameter x for storage purposes as

_to_bfloat16_unbiased(tf.pow(x, 5)) . This gives 5x the

resolution while incurring overflow and underflow at 10^9 and 10^-9

instead of 10^37 and 10^-37. Comes at a cost of extracting fifth roots

to decode parameters. Or use some other such scheme.

2. In this function, use actual random numbers, different for each parameter

as opposed to the same for every parameter in the graph.

3. Look for bugs in this function.

Args:

x: A float32 Tensor.

Returns:

A float32 Tensor.

"""

# Not using random_uniform here due to a problem on TPU in that random seeds

# are not respected, which may cause the parameters on different replicas

# to go out-of-sync.

quantization_noise = _quantization_noise_from_step_num()

x_sign = tf.sign(x)

# Make sure x is positive. If it is zero, the two candidates are identical.

x = x * x_sign + 1e-30

cand1 = tf.to_bfloat16(x)

cand1_f = tf.to_float(cand1)

# This relies on the fact that for a positive bfloat16 b,

# b * 1.005 gives you the next higher bfloat16 and b*0.995 gives you the

# next lower one. Both 1.005 and 0.995 are ballpark estimation.

cand2 = tf.to_bfloat16(

tf.where(tf.greater(x, cand1_f), cand1_f * 1.005, cand1_f * 0.995))

ret = _randomized_roundoff_to_bfloat16(x, quantization_noise, cand1, cand2)

return ret * tf.to_bfloat16(x_sign)