t5x/losses.py

# Copyright 2022 The T5X 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.

"""Loss functions."""
import enum
from typing import Tuple, Mapping, Optional, Union

from flax.training import common_utils
import jax
import jax.numpy as jnp
import numpy as np


@jax.custom_vjp
def cross_entropy_with_logits(logits: jnp.ndarray, targets: jnp.ndarray,
                              z_loss: float) -> jnp.ndarray:
  """Computes cross entropy loss with stable custom gradient.

  Computes a stabilized-gradient version of:
    -jnp.sum(targets * nn.log_softmax(logits), axis=-1)

  If z_loss > 0, then an auxiliary loss equal to z_loss*log(z)^2
  will be added to the cross entropy loss (z = softmax normalization constant).
  The two uses of z_loss are:
  1. To keep the logits from drifting too far from zero, which can cause
     unacceptable roundoff errors in bfloat16.
  2. To encourage the logits to be normalized log-probabilities.

  Args:
    logits: [batch, length, num_classes] float array.
    targets: categorical one-hot targets [batch, length, num_classes] float
      array.
    z_loss: coefficient for auxilliary z-loss loss term.

  Returns:
    tuple with the total loss and the z_loss, both
    float arrays with shape [batch, length].
  """
  logits_sum = jax.scipy.special.logsumexp(logits, axis=-1, keepdims=True)
  log_softmax = logits - logits_sum
  loss = -jnp.sum(targets * log_softmax, axis=-1)
  # Add auxilliary z-loss term.
  log_z = jnp.squeeze(logits_sum, axis=-1)
  total_z_loss = z_loss * jax.lax.square(log_z)
  loss += total_z_loss
  return loss, total_z_loss


def _cross_entropy_with_logits_fwd(
    logits: jnp.ndarray,
    targets: jnp.ndarray,
    z_loss: float = 0.0
) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp
                              .ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]]:
  """Forward-mode of `cross_entropy_with_logits`."""
  max_logit = logits.max(axis=-1, keepdims=True)
  shifted = logits - max_logit
  exp_shifted = jnp.exp(shifted)
  sum_exp = jnp.sum(exp_shifted, axis=-1, keepdims=True)
  log_softmax = shifted - jnp.log(sum_exp)
  loss = -jnp.sum(targets * log_softmax, axis=-1)
  # Add auxilliary z-loss term.
  log_z = jnp.squeeze(jnp.log(sum_exp) + max_logit, axis=-1)
  total_z_loss = z_loss * jax.lax.square(log_z)
  loss += total_z_loss
  return (loss, total_z_loss), (logits, targets, z_loss, exp_shifted, sum_exp,
                                log_softmax, log_z)


def _cross_entropy_with_logits_bwd(
    res: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray,
               jnp.ndarray, jnp.ndarray], g: Tuple[jnp.ndarray, jnp.ndarray]
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
  """Backward-mode of `cross_entropy_with_logits`."""
  g = g[0]  # Ignore z_loss component as that is only used for logging.
  logits, targets, z_loss, exp_shifted, sum_exp, log_softmax, log_z = res
  # z-loss term adds the (2 * z_loss * log_z) factor.
  deriv = (
      jnp.expand_dims(1 + 2 * z_loss * log_z, -1) * exp_shifted / sum_exp -
      targets)
  g_logits = jnp.expand_dims(g, axis=-1) * deriv
  g_targets = -jnp.expand_dims(g, axis=-1) * log_softmax
  return (jnp.asarray(g_logits,
                      logits.dtype), jnp.asarray(g_targets, targets.dtype),
          jnp.array(0.0))  # sets z-loss coeff gradient to 0


cross_entropy_with_logits.defvjp(_cross_entropy_with_logits_fwd,
                                 _cross_entropy_with_logits_bwd)


def compute_weighted_cross_entropy(
    logits: jnp.ndarray,
    targets: jnp.ndarray,
    weights: Optional[jnp.ndarray] = None,
    label_smoothing: float = 0.0,
    z_loss: float = 0.0,
    loss_normalizing_factor: Optional[float] = None
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
  """Compute weighted cross entropy and entropy for log probs and targets.

  Args:
   logits: [batch, length, num_classes] float array.
   targets: categorical targets [batch, length] int array.
   weights: None or array of shape [batch, length].
   label_smoothing: label smoothing constant, used to determine the on and off
     values.
   z_loss: coefficient for auxiliary z-loss loss term.
   loss_normalizing_factor: Constant to divide loss by. If not specified, loss
     will not be normalized. Intended for backward compatibility with T5-MTF
     training. Should not normally be used.

  Returns:
    Tuple of scalar loss, z_loss, and weight sum.
  """
  if logits.ndim != targets.ndim + 1:
    raise ValueError('Incorrect shapes. Got shape %s logits and %s targets' %
                     (str(logits.shape), str(targets.shape)))
  vocab_size = logits.shape[-1]
  confidence = 1.0 - label_smoothing
  low_confidence = (1.0 - confidence) / (vocab_size - 1)
  normalizing_constant = -(
      confidence * jnp.log(confidence) +
      (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20))
  soft_targets = common_utils.onehot(
      targets, vocab_size, on_value=confidence, off_value=low_confidence)
  total_loss, total_z_loss = cross_entropy_with_logits(
      logits, soft_targets, z_loss=z_loss)
  total_loss = total_loss - normalizing_constant

  weight_sum = np.prod(targets.shape)
  if weights is not None:
    total_loss = total_loss * weights
    total_z_loss = total_z_loss * weights
    weight_sum = jnp.sum(weights)

  # By default, we do not normalize loss based on anything.
  # We don't normalize based on batch size because the optimizers we use are
  # pretty much scale invariant, so this simplifies things.
  # We don't normalize based on number of non-padding tokens in order to treat
  # each token as equally important regardless of sequence length.
  if loss_normalizing_factor is not None:
    total_loss /= loss_normalizing_factor
    total_z_loss /= loss_normalizing_factor
  return jnp.sum(total_loss), jnp.sum(total_z_loss), weight_sum


@enum.unique
class SpecialLossNormalizingFactor(enum.Enum):
  """Specially calcualted loss_normalizing_factors, that are not a constant.

  Attributes:
    NUM_REAL_TARGET_TOKENS: Whether to divide the loss by the number of real
      (non-padding) tokens in the current target batch. If
      'decoder_loss_weights' are specified, it will be the sum of the weights.
      Otherwise it will be the number of non-zero 'decoder_target_tokens'.
    NUM_TOTAL_TARGET_TOKENS: Whether to divide the loss by the total number of
      target tokens, i.e., batch_size * target_seq_length (including padding).
    AVERAGE_PER_SEQUENCE: This will first compute the per-sequence loss
      (averaged over the number of real target tokens in the sequence), and then
      compute the average of that over the sequences. This can be preferable to
      NUM_REAL_TARGET_TOKENS for finetuning, because it will weigh all examples
      equally, regardless of sequence length (which can be especially important
      for multi-task finetuning).
  """
  NUM_REAL_TARGET_TOKENS = 1
  NUM_TOTAL_TARGET_TOKENS = 2
  AVERAGE_PER_SEQUENCE = 3


def convert_special_loss_normalizing_factor_to_enum(
    x: str) -> SpecialLossNormalizingFactor:
  """Converts stringified version of LNF to an enum.

  This is useful because gin dynamic registration does not (currently)
  have support for enum.

  Args:
    x: stringified version of SpecialLossNormalizingFactor enum.

  Returns:
    SpecialLossNormalizingFactor enum instance.
  """
  x = x.upper()
  if x == 'NUM_REAL_TARGET_TOKENS':
    return SpecialLossNormalizingFactor.NUM_REAL_TARGET_TOKENS
  if x == 'NUM_TOTAL_TARGET_TOKENS':
    return SpecialLossNormalizingFactor.NUM_TOTAL_TARGET_TOKENS
  if x == 'AVERAGE_PER_SEQUENCE':
    return SpecialLossNormalizingFactor.AVERAGE_PER_SEQUENCE
  raise ValueError(
      'Could not convert string \"%s\" to SpecialLossNormalizingFactor' % x)


def get_loss_normalizing_factor_and_weights(
    loss_normalizing_factor: Optional[Union[float, int, str,
                                            SpecialLossNormalizingFactor]],
    batch: Mapping[str, jnp.ndarray]):
  """Get the float loss_normalizing_factor and loss weights.

  If loss_normalizing_factor is float or None, this will simply return the
  input loss_normalizing_factor and batch.

  If loss_normalizing_factor is a SpecialLossNormalizingFactor, it will
  return a float loss_normalizing_factor and loss weights corresponding to
  the special LNF. See SpecialLossNormalizingFactor for more details.

  Args:
    loss_normalizing_factor: The input LNF, which may be a float, None, or
      SpecialLossNormalizingFactor (or a stringified SLNF).
    batch: Input data batch.

  Returns:
    Tuple of (output_loss_normalizing_factor, loss_weights).
      'output_loss_normalizing_factor' is a scalar float (Python float
      or jnp float).
      'loss_weights' is the per token loss weight JNP array.
  """

  loss_weights = batch.get('decoder_loss_weights', None)
  if (loss_normalizing_factor is None or
      not isinstance(loss_normalizing_factor,
                     (str, SpecialLossNormalizingFactor))):
    return (loss_normalizing_factor, loss_weights)

  if isinstance(loss_normalizing_factor, str):
    loss_normalizing_factor = convert_special_loss_normalizing_factor_to_enum(
        loss_normalizing_factor)

  # If `loss_weights` are not provided, we assume that the padding id is 0 and
  # that non-padding tokens in the decoder all correspond to the positions
  # where loss should be taken. If more fine-grained behavior (e.g., taking
  # loss on subset of 'decoder_target_tokens') is desired, provide
  # `loss_weights` that account for this.
  if loss_weights is None:
    loss_weights = jnp.asarray(batch['decoder_target_tokens'] > 0, jnp.float32)

  output_normalizing_factor = None
  if (loss_normalizing_factor ==
      SpecialLossNormalizingFactor.NUM_REAL_TARGET_TOKENS):
    output_normalizing_factor = jnp.sum(loss_weights)
  elif (loss_normalizing_factor ==
        SpecialLossNormalizingFactor.NUM_TOTAL_TARGET_TOKENS):
    output_normalizing_factor = np.prod(batch['decoder_target_tokens'].shape)
  elif (loss_normalizing_factor ==
        SpecialLossNormalizingFactor.AVERAGE_PER_SEQUENCE):
    loss_weights /= jnp.sum(loss_weights, axis=-1, keepdims=True) + 1e-3
    output_normalizing_factor = jnp.sum(loss_weights)
  else:
    raise ValueError('Unsupported value of loss_normalizing_factor: %s' %
                     str(loss_normalizing_factor))

  return (output_normalizing_factor, loss_weights)