<reponame>brain-research/fisher-rao-regularization<gh_stars>1-10
# Copyright 2018 Google LLC
#
# 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.

"""Finite difference approximations of the Fisher-Rao norm regularizer.

The provided utility routines are used to create Fisher-Rao norm regularizers.
The implementations use finite difference perturbations of the parameters in the
original loss to approximate the necessary gradient-vector products.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf


class VariableCollector(object):
  """Helper class with custom getter to collect `Variable` objects in a scope.

  When called for the first time the custom getter stores the corresponding
  `Variable` object in dictionary.
  When called subsequently for the same `Variable` name the getter will return
  the object from the dictionary instead of calling the original getter.
  """

  def __init__(self):
    self.variables = {}

  def collector_getter(self, getter, name, *args, **kwargs):
    """Custom getter for `VariableScope` that stores `Variable` in dictionary.

    Args:
      getter: Function, original getter function
      name: String, name of `Variable` in the scope
      *args: Additional arguments, currently only passed forward on first
        call of the original getter
      **kwargs: Additional arguments, currently only passed forward on first
        call of the original getter

    Returns:
      A `Tensor` object that contains the named `Variable` either from calling
      the original getter or if available from the dictionary.
    """

    if name not in self.variables:
      self.variables[name] = getter(name, *args, **kwargs)

    # TODO(jonathanjh): Add consistency check for args and kwargs.

    return self.variables[name]


def make_perturbation_getter(should_regularize, collector, perturbation):
  """Creates custom getter to replace variables in scope by their perturbations.

  Args:
    should_regularize: Function, takes a variable name as String and returns
      Boolean that decides whether the variable should be regularized.
    collector: `VariableCollector` object that provides the dictionary to use
      for subsequent use of the same `Variable` object for the same name.
    perturbation: Float, perturbation value added on top of 1 to use for
      variable replacement value.

  Returns:
    A custom getter function that can be used with `VariableScope`.
  """

  def plus_getter(getter, name, *args, **kwargs):
    var = collector.collector_getter(getter, name, *args, **kwargs)

    if should_regularize(name) and kwargs.get("trainable"):
      var = (1. + perturbation) * var

    return var

  return plus_getter


def make_empirical_fisher_regularizer(make_logits, labels, scope,
                                      should_regularize, perturbation):
  """Creates per-example logits and the per-example empirical Fisher-Rao norm.

  This function assumes the model of a categorical distribution generated by a
  softmax function.
  The empirical Fisher-Rao norm uses the empirical training distribution for
  both the input values and the labels to estimate the Fisher information
  matrix.

  Args:
    make_logits: Function, returns `Tensor` representing the per-example logits.
      The expected shape of the tensor is such that the number of categories
      is the last dimension.
    labels: Tensor, encoding of the class labels compatible in dimension with
      the return of the make_logits function.
    scope: String, name of `VariableScope` to use for the `Variable` objects
      that represent the regularized parameter.
    should_regularize: Function, takes a variable name as String and returns
      Boolean that decides whether the variable should be regularized.
      The passed variable name includes the name of the scope.
    perturbation: Float, finite difference perturbation constant.
      The choice of perturbation constant represents a tradeoff between rounding
      and approximation error and should depend on floating point precision and
      parameter norm.

  Returns:
    A tuple of `Tensor` objects representing the per-example logits and the
    scalar empirical Fisher-Rao norm regularization loss.

  Raises:
    ValueError: if the last dimension of the logits shape is not statically
    inferrable.
  """

  collector = VariableCollector()

  with tf.variable_scope(scope, custom_getter=collector.collector_getter):
    logits = make_logits()

  if logits.shape[-1].value is None:
    raise ValueError("The size of the last dimension of the logits vector must"
                     " be statically inferrable.")

  with tf.variable_scope(
      scope,
      custom_getter=
      make_perturbation_getter(should_regularize, collector, perturbation)):
    perturbed_logits = make_logits()

  loss = tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits)
  perturbed_loss = tf.nn.softmax_cross_entropy_with_logits(
      labels=labels, logits=perturbed_logits)

  regularizer = tf.square(
      tf.divide(tf.subtract(perturbed_loss, loss), perturbation))

  regularizer = tf.reduce_mean(regularizer)

  return (logits, regularizer)


def make_standard_fisher_regularizer(make_logits, scope, should_regularize,
                                     perturbation, differentiate_probability):
  """Creates per-example logits and the per-example standard Fisher-Rao norm.

  This function assumes the model of a categorical distribution generated by a
  softmax function.
  The standard Fisher-Rao norm uses the model distribution computed from the
  logits by the softmax function to estimate the Fisher information matrix.
  The empirical training distribution is used for the input values.

  Args:
    make_logits: Function, returns `Tensor` representing the per-example logits.
      The expected shape of the tensor is such that the number of categories
      is the last dimension.
    scope: String, name of `VariableScope` to use for the `Variable` objects
      that represent the regularized parameter.
    should_regularize: Function, takes a variable name as String and returns
      Boolean that decides whether the variable should be regularized.
      The passed variable name includes the name of the scope.
    perturbation: Float, finite difference perturbation constant.
      The choice of perturbation constant represents a tradeoff between rounding
      and approximation error and should depend on floating point precision and
      parameter norm.
    differentiate_probability: Boolean, determines whether the label probability
      distribution should be differentiated.

  Returns:
    A tuple of `Tensor` objects representing the per-example logits and the
    scalar standard Fisher-Rao norm regularization loss.

  Raises:
    ValueError: if the last dimension of the logits shape is not statically
    inferrable.
  """

  collector = VariableCollector()

  with tf.variable_scope(scope, custom_getter=collector.collector_getter):
    logits = make_logits()

  if logits.shape[-1].value is None:
    raise ValueError("The size of the last dimension of the logits vector must"
                     " be statically inferrable.")

  with tf.variable_scope(
      scope,
      custom_getter=
      make_perturbation_getter(should_regularize, collector, perturbation)):
    perturbed_logits = make_logits()

  log_probs = tf.nn.log_softmax(logits, axis=-1)
  perturbed_log_probs = tf.nn.log_softmax(perturbed_logits, axis=-1)

  stop_probs = tf.stop_gradient(tf.exp(log_probs))
  log_prob_derivative = (tf.square((perturbed_log_probs - log_probs) /
                                   perturbation))

  if differentiate_probability:
    prob_regularizer_loss = (log_prob_derivative * stop_probs +
                             tf.stop_gradient(log_prob_derivative) * log_probs *
                             stop_probs -
                             tf.stop_gradient(log_prob_derivative * log_probs *
                                              stop_probs))

  else:
    prob_regularizer_loss = log_prob_derivative * stop_probs

  regularizer = logits.shape[-1].value * tf.reduce_mean(prob_regularizer_loss)

  return (logits, regularizer)


# Copyright 2018 Google LLC
#
# 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.

"""Finite difference approximations of the Fisher-Rao norm regularizer.

The provided utility routines are used to create Fisher-Rao norm regularizers.
The implementations use finite difference perturbations of the parameters in the
original loss to approximate the necessary gradient-vector products.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf


class VariableCollector(object):
  """Helper class with custom getter to collect `Variable` objects in a scope.

  When called for the first time the custom getter stores the corresponding
  `Variable` object in dictionary.
  When called subsequently for the same `Variable` name the getter will return
  the object from the dictionary instead of calling the original getter.
  """

  def __init__(self):
    self.variables = {}

  def collector_getter(self, getter, name, *args, **kwargs):
    """Custom getter for `VariableScope` that stores `Variable` in dictionary.

    Args:
      getter: Function, original getter function
      name: String, name of `Variable` in the scope
      *args: Additional arguments, currently only passed forward on first
        call of the original getter
      **kwargs: Additional arguments, currently only passed forward on first
        call of the original getter

    Returns:
      A `Tensor` object that contains the named `Variable` either from calling
      the original getter or if available from the dictionary.
    """

    if name not in self.variables:
      self.variables[name] = getter(name, *args, **kwargs)

    # TODO(jonathanjh): Add consistency check for args and kwargs.

    return self.variables[name]


def make_perturbation_getter(should_regularize, collector, perturbation):
  """Creates custom getter to replace variables in scope by their perturbations.

  Args:
    should_regularize: Function, takes a variable name as String and returns
      Boolean that decides whether the variable should be regularized.
    collector: `VariableCollector` object that provides the dictionary to use
      for subsequent use of the same `Variable` object for the same name.
    perturbation: Float, perturbation value added on top of 1 to use for
      variable replacement value.

  Returns:
    A custom getter function that can be used with `VariableScope`.
  """

  def plus_getter(getter, name, *args, **kwargs):
    var = collector.collector_getter(getter, name, *args, **kwargs)

    if should_regularize(name) and kwargs.get("trainable"):
      var = (1. + perturbation) * var

    return var

  return plus_getter


def make_empirical_fisher_regularizer(make_logits, labels, scope,
                                      should_regularize, perturbation):
  """Creates per-example logits and the per-example empirical Fisher-Rao norm.

  This function assumes the model of a categorical distribution generated by a
  softmax function.
  The empirical Fisher-Rao norm uses the empirical training distribution for
  both the input values and the labels to estimate the Fisher information
  matrix.

  Args:
    make_logits: Function, returns `Tensor` representing the per-example logits.
      The expected shape of the tensor is such that the number of categories
      is the last dimension.
    labels: Tensor, encoding of the class labels compatible in dimension with
      the return of the make_logits function.
    scope: String, name of `VariableScope` to use for the `Variable` objects
      that represent the regularized parameter.
    should_regularize: Function, takes a variable name as String and returns
      Boolean that decides whether the variable should be regularized.
      The passed variable name includes the name of the scope.
    perturbation: Float, finite difference perturbation constant.
      The choice of perturbation constant represents a tradeoff between rounding
      and approximation error and should depend on floating point precision and
      parameter norm.

  Returns:
    A tuple of `Tensor` objects representing the per-example logits and the
    scalar empirical Fisher-Rao norm regularization loss.

  Raises:
    ValueError: if the last dimension of the logits shape is not statically
    inferrable.
  """

  collector = VariableCollector()

  with tf.variable_scope(scope, custom_getter=collector.collector_getter):
    logits = make_logits()

  if logits.shape[-1].value is None:
    raise ValueError("The size of the last dimension of the logits vector must"
                     " be statically inferrable.")

  with tf.variable_scope(
      scope,
      custom_getter=
      make_perturbation_getter(should_regularize, collector, perturbation)):
    perturbed_logits = make_logits()

  loss = tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits)
  perturbed_loss = tf.nn.softmax_cross_entropy_with_logits(
      labels=labels, logits=perturbed_logits)

  regularizer = tf.square(
      tf.divide(tf.subtract(perturbed_loss, loss), perturbation))

  regularizer = tf.reduce_mean(regularizer)

  return (logits, regularizer)


def make_standard_fisher_regularizer(make_logits, scope, should_regularize,
                                     perturbation, differentiate_probability):
  """Creates per-example logits and the per-example standard Fisher-Rao norm.

  This function assumes the model of a categorical distribution generated by a
  softmax function.
  The standard Fisher-Rao norm uses the model distribution computed from the
  logits by the softmax function to estimate the Fisher information matrix.
  The empirical training distribution is used for the input values.

  Args:
    make_logits: Function, returns `Tensor` representing the per-example logits.
      The expected shape of the tensor is such that the number of categories
      is the last dimension.
    scope: String, name of `VariableScope` to use for the `Variable` objects
      that represent the regularized parameter.
    should_regularize: Function, takes a variable name as String and returns
      Boolean that decides whether the variable should be regularized.
      The passed variable name includes the name of the scope.
    perturbation: Float, finite difference perturbation constant.
      The choice of perturbation constant represents a tradeoff between rounding
      and approximation error and should depend on floating point precision and
      parameter norm.
    differentiate_probability: Boolean, determines whether the label probability
      distribution should be differentiated.

  Returns:
    A tuple of `Tensor` objects representing the per-example logits and the
    scalar standard Fisher-Rao norm regularization loss.

  Raises:
    ValueError: if the last dimension of the logits shape is not statically
    inferrable.
  """

  collector = VariableCollector()

  with tf.variable_scope(scope, custom_getter=collector.collector_getter):
    logits = make_logits()

  if logits.shape[-1].value is None:
    raise ValueError("The size of the last dimension of the logits vector must"
                     " be statically inferrable.")

  with tf.variable_scope(
      scope,
      custom_getter=
      make_perturbation_getter(should_regularize, collector, perturbation)):
    perturbed_logits = make_logits()

  log_probs = tf.nn.log_softmax(logits, axis=-1)
  perturbed_log_probs = tf.nn.log_softmax(perturbed_logits, axis=-1)

  stop_probs = tf.stop_gradient(tf.exp(log_probs))
  log_prob_derivative = (tf.square((perturbed_log_probs - log_probs) /
                                   perturbation))

  if differentiate_probability:
    prob_regularizer_loss = (log_prob_derivative * stop_probs +
                             tf.stop_gradient(log_prob_derivative) * log_probs *
                             stop_probs -
                             tf.stop_gradient(log_prob_derivative * log_probs *
                                              stop_probs))

  else:
    prob_regularizer_loss = log_prob_derivative * stop_probs

  regularizer = logits.shape[-1].value * tf.reduce_mean(prob_regularizer_loss)

  return (logits, regularizer)
# Copyright 2018 Google LLC # # 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. Finite difference approximations of the Fisher-Rao norm regularizer.  The provided utility routines are used to create Fisher-Rao norm regularizers. The implementations use finite difference perturbations of the parameters in the original loss to approximate the necessary gradient-vector products. Helper class with custom getter to collect `Variable` objects in a scope.    When called for the first time the custom getter stores the corresponding   `Variable` object in dictionary.   When called subsequently for the same `Variable` name the getter will return   the object from the dictionary instead of calling the original getter. Custom getter for `VariableScope` that stores `Variable` in dictionary.      Args:       getter: Function, original getter function       name: String, name of `Variable` in the scope       *args: Additional arguments, currently only passed forward on first         call of the original getter       **kwargs: Additional arguments, currently only passed forward on first         call of the original getter      Returns:       A `Tensor` object that contains the named `Variable` either from calling       the original getter or if available from the dictionary. # TODO(jonathanjh): Add consistency check for args and kwargs. Creates custom getter to replace variables in scope by their perturbations.    Args:     should_regularize: Function, takes a variable name as String and returns       Boolean that decides whether the variable should be regularized.     collector: `VariableCollector` object that provides the dictionary to use       for subsequent use of the same `Variable` object for the same name.     perturbation: Float, perturbation value added on top of 1 to use for       variable replacement value.    Returns:     A custom getter function that can be used with `VariableScope`. Creates per-example logits and the per-example empirical Fisher-Rao norm.    This function assumes the model of a categorical distribution generated by a   softmax function.   The empirical Fisher-Rao norm uses the empirical training distribution for   both the input values and the labels to estimate the Fisher information   matrix.    Args:     make_logits: Function, returns `Tensor` representing the per-example logits.       The expected shape of the tensor is such that the number of categories       is the last dimension.     labels: Tensor, encoding of the class labels compatible in dimension with       the return of the make_logits function.     scope: String, name of `VariableScope` to use for the `Variable` objects       that represent the regularized parameter.     should_regularize: Function, takes a variable name as String and returns       Boolean that decides whether the variable should be regularized.       The passed variable name includes the name of the scope.     perturbation: Float, finite difference perturbation constant.       The choice of perturbation constant represents a tradeoff between rounding       and approximation error and should depend on floating point precision and       parameter norm.    Returns:     A tuple of `Tensor` objects representing the per-example logits and the     scalar empirical Fisher-Rao norm regularization loss.    Raises:     ValueError: if the last dimension of the logits shape is not statically     inferrable. Creates per-example logits and the per-example standard Fisher-Rao norm.    This function assumes the model of a categorical distribution generated by a   softmax function.   The standard Fisher-Rao norm uses the model distribution computed from the   logits by the softmax function to estimate the Fisher information matrix.   The empirical training distribution is used for the input values.    Args:     make_logits: Function, returns `Tensor` representing the per-example logits.       The expected shape of the tensor is such that the number of categories       is the last dimension.     scope: String, name of `VariableScope` to use for the `Variable` objects       that represent the regularized parameter.     should_regularize: Function, takes a variable name as String and returns       Boolean that decides whether the variable should be regularized.       The passed variable name includes the name of the scope.     perturbation: Float, finite difference perturbation constant.       The choice of perturbation constant represents a tradeoff between rounding       and approximation error and should depend on floating point precision and       parameter norm.     differentiate_probability: Boolean, determines whether the label probability       distribution should be differentiated.    Returns:     A tuple of `Tensor` objects representing the per-example logits and the     scalar standard Fisher-Rao norm regularization loss.    Raises:     ValueError: if the last dimension of the logits shape is not statically     inferrable.