google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Defines common utilties for variational dropout layers."""
17
from __future__ import absolute_import
18
from __future__ import division
19
from __future__ import print_function
20

21
import tensorflow.compat.v1 as tf
22

23

24
EPSILON = 1e-8
25

26

27
def compute_log_alpha(log_sigma2, theta, eps=EPSILON, value_limit=8.):
28
  R"""Compute the log \alpha values from \theta and log \sigma^2.
29

30
  The relationship between \sigma^2, \theta, and \alpha as defined in the
31
  paper https://arxiv.org/abs/1701.05369 is
32

33
  \sigma^2 = \alpha * \theta^2
34

35
  This method calculates the log \alpha values based on this relation.
36

37
  Args:
38
    log_sigma2: tf.Variable. The log variance for each weight.
39
    theta: tf.Variable. The mean for each weight.
40
    eps: Small constant value to use in log and sqrt operations to avoid NaNs.
41
    value_limit: If not None, the log_alpha values will be clipped to the
42
     range [-value_limit, value_limit]. This is consistent with the
43
     implementation provided with the publication.
44

45
  Returns:
46
    A tf.Tensor representing the calculated log \alpha values.
47
  """
48
  log_alpha = log_sigma2 - tf.log(tf.square(theta) + eps)
49

50
  if value_limit is not None:
51
    # If a limit is specified, clip the alpha values
52
    return tf.clip_by_value(log_alpha, -value_limit, value_limit)
53
  return log_alpha
54

55

56
def compute_log_sigma2(log_alpha, theta, eps=EPSILON):
57
  R"""Compute the log \sigma^2 values from log \alpha and \theta.
58

59
  The relationship between \sigma^2, \theta, and \alpha as defined in the
60
  paper https://arxiv.org/abs/1701.05369 is
61

62
  \sigma^2 = \alpha * \theta^2
63

64
  This method calculates the log \sigma^2 values based on this relation.
65

66
  Args:
67
    log_alpha: tf.Tensor. The log alpha values for each weight.
68
    theta: tf.Variable. The mean for each weight.
69
    eps: Small constant value to use in log and sqrt operations to avoid NaNs.
70

71
  Returns:
72
    A tf.Tensor representing the calculated log \sigma^2 values.
73
  """
74
  return log_alpha + tf.log(tf.square(theta) + eps)
75