# Copyright 2023 The TensorFlow Authors. All Rights Reserved. # # 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. """Definitions for random feature Gaussian process layer.""" import math import tensorflow as tf, tf_keras _SUPPORTED_LIKELIHOOD = ('binary_logistic', 'poisson', 'gaussian') class RandomFeatureGaussianProcess(tf_keras.layers.Layer): """Gaussian process layer with random feature approximation [1]. During training, the model updates the maximum a posteriori (MAP) logits estimates and posterior precision matrix using minibatch statistics. During inference, the model divides the MAP logit estimates by the predictive standard deviation, which is equivalent to approximating the posterior mean of the predictive probability via the mean-field approximation. User can specify different types of random features by setting `use_custom_random_features=True`, and change the initializer and activations of the custom random features. For example: MLP Kernel: initializer='random_normal', activation=tf.nn.relu RBF Kernel: initializer='random_normal', activation=tf.math.cos A linear kernel can also be specified by setting gp_kernel_type='linear' and `use_custom_random_features=True`. [1]: Ali Rahimi and Benjamin Recht. Random Features for Large-Scale Kernel Machines. In _Neural Information Processing Systems_, 2007. https://people.eecs.berkeley.edu/~brecht/papers/07.rah.rec.nips.pdf Attributes: units: (int) The dimensionality of layer. num_inducing: (int) The number of random features for the approximation. is_training: (tf.bool) Whether the layer is set in training mode. If so the layer updates the Gaussian process' variance estimate using statistics computed from the incoming minibatches. """ def __init__(self, units, num_inducing=1024, gp_kernel_type='gaussian', gp_kernel_scale=1., gp_output_bias=0., normalize_input=False, gp_kernel_scale_trainable=False, gp_output_bias_trainable=False, gp_cov_momentum=0.999, gp_cov_ridge_penalty=1., scale_random_features=True, use_custom_random_features=True, custom_random_features_initializer=None, custom_random_features_activation=None, l2_regularization=1e-6, gp_cov_likelihood='gaussian', return_gp_cov=True, return_random_features=False, dtype=None, name='random_feature_gaussian_process', **gp_output_kwargs): """Initializes a random-feature Gaussian process layer instance. Args: units: (int) Number of output units. num_inducing: (int) Number of random Fourier features used for approximating the Gaussian process. gp_kernel_type: (string) The type of kernel function to use for Gaussian process. Currently default to 'gaussian' which is the Gaussian RBF kernel. gp_kernel_scale: (float) The length-scale parameter of the a shift-invariant kernel function, i.e., for RBF kernel: exp(-|x1 - x2|**2 / gp_kernel_scale). gp_output_bias: (float) Scalar initial value for the bias vector. normalize_input: (bool) Whether to normalize the input to Gaussian process. gp_kernel_scale_trainable: (bool) Whether the length scale variable is trainable. gp_output_bias_trainable: (bool) Whether the bias is trainable. gp_cov_momentum: (float) A discount factor used to compute the moving average for posterior covariance matrix. gp_cov_ridge_penalty: (float) Initial Ridge penalty to posterior covariance matrix. scale_random_features: (bool) Whether to scale the random feature by sqrt(2. / num_inducing). use_custom_random_features: (bool) Whether to use custom random features implemented using tf_keras.layers.Dense. custom_random_features_initializer: (str or callable) Initializer for the random features. Default to random normal which approximates a RBF kernel function if activation function is cos. custom_random_features_activation: (callable) Activation function for the random feature layer. Default to cosine which approximates a RBF kernel function. l2_regularization: (float) The strength of l2 regularization on the output weights. gp_cov_likelihood: (string) Likelihood to use for computing Laplace approximation for covariance matrix. Default to `gaussian`. return_gp_cov: (bool) Whether to also return GP covariance matrix. If False then no covariance learning is performed. return_random_features: (bool) Whether to also return random features. dtype: (tf.DType) Input data type. name: (string) Layer name. **gp_output_kwargs: Additional keyword arguments to dense output layer. """ super().__init__(name=name, dtype=dtype) self.units = units self.num_inducing = num_inducing self.normalize_input = normalize_input self.gp_input_scale = 1. / tf.sqrt(gp_kernel_scale) self.gp_feature_scale = tf.sqrt(2. / float(num_inducing)) self.scale_random_features = scale_random_features self.return_random_features = return_random_features self.return_gp_cov = return_gp_cov self.gp_kernel_type = gp_kernel_type self.gp_kernel_scale = gp_kernel_scale self.gp_output_bias = gp_output_bias self.gp_kernel_scale_trainable = gp_kernel_scale_trainable self.gp_output_bias_trainable = gp_output_bias_trainable self.use_custom_random_features = use_custom_random_features self.custom_random_features_initializer = custom_random_features_initializer self.custom_random_features_activation = custom_random_features_activation self.l2_regularization = l2_regularization self.gp_output_kwargs = gp_output_kwargs self.gp_cov_momentum = gp_cov_momentum self.gp_cov_ridge_penalty = gp_cov_ridge_penalty self.gp_cov_likelihood = gp_cov_likelihood if self.use_custom_random_features: # Default to Gaussian RBF kernel. self.random_features_bias_initializer = tf.random_uniform_initializer( minval=0., maxval=2. * math.pi) if self.custom_random_features_initializer is None: self.custom_random_features_initializer = ( tf_keras.initializers.RandomNormal(stddev=1.)) if self.custom_random_features_activation is None: self.custom_random_features_activation = tf.math.cos def build(self, input_shape): # Defines model layers. if self.normalize_input: self._input_norm_layer = tf_keras.layers.LayerNormalization( name='gp_input_normalization') self._input_norm_layer.build(input_shape) input_shape = self._input_norm_layer.compute_output_shape(input_shape) self._random_feature = self._make_random_feature_layer( name='gp_random_feature') self._random_feature.build(input_shape) input_shape = self._random_feature.compute_output_shape(input_shape) if self.return_gp_cov: self._gp_cov_layer = LaplaceRandomFeatureCovariance( momentum=self.gp_cov_momentum, ridge_penalty=self.gp_cov_ridge_penalty, likelihood=self.gp_cov_likelihood, dtype=self.dtype, name='gp_covariance') self._gp_cov_layer.build(input_shape) self._gp_output_layer = tf_keras.layers.Dense( units=self.units, use_bias=False, kernel_regularizer=tf_keras.regularizers.l2(self.l2_regularization), dtype=self.dtype, name='gp_output_weights', **self.gp_output_kwargs) self._gp_output_layer.build(input_shape) self._gp_output_bias = tf.Variable( initial_value=[self.gp_output_bias] * self.units, dtype=self.dtype, trainable=self.gp_output_bias_trainable, name='gp_output_bias') self.built = True def _make_random_feature_layer(self, name): """Defines random feature layer depending on kernel type.""" if not self.use_custom_random_features: # Use default RandomFourierFeatures layer from tf_keras. return tf_keras.layers.experimental.RandomFourierFeatures( output_dim=self.num_inducing, kernel_initializer=self.gp_kernel_type, scale=self.gp_kernel_scale, trainable=self.gp_kernel_scale_trainable, dtype=self.dtype, name=name) if self.gp_kernel_type.lower() == 'linear': custom_random_feature_layer = tf_keras.layers.Lambda( lambda x: x, name=name) else: # Use user-supplied configurations. custom_random_feature_layer = tf_keras.layers.Dense( units=self.num_inducing, use_bias=True, activation=self.custom_random_features_activation, kernel_initializer=self.custom_random_features_initializer, bias_initializer=self.random_features_bias_initializer, trainable=False, name=name) return custom_random_feature_layer def reset_covariance_matrix(self): """Resets covariance matrix of the GP layer. This function is useful for reseting the model's covariance matrix at the beginning of a new epoch. """ self._gp_cov_layer.reset_precision_matrix() def call(self, inputs, global_step=None, training=None): # Computes random features. gp_inputs = inputs if self.normalize_input: gp_inputs = self._input_norm_layer(gp_inputs) elif self.use_custom_random_features: # Supports lengthscale for custom random feature layer by directly # rescaling the input. gp_input_scale = tf.cast(self.gp_input_scale, inputs.dtype) gp_inputs = gp_inputs * gp_input_scale gp_feature = self._random_feature(gp_inputs) if self.scale_random_features: # Scale random feature by 2. / sqrt(num_inducing) following [1]. # When using GP layer as the output layer of a nerual network, # it is recommended to turn this scaling off to prevent it from changing # the learning rate to the hidden layers. gp_feature_scale = tf.cast(self.gp_feature_scale, inputs.dtype) gp_feature = gp_feature * gp_feature_scale # Computes posterior center (i.e., MAP estimate) and variance. gp_output = self._gp_output_layer(gp_feature) + self._gp_output_bias if self.return_gp_cov: gp_covmat = self._gp_cov_layer(gp_feature, gp_output, training) # Assembles model output. model_output = [gp_output,] if self.return_gp_cov: model_output.append(gp_covmat) if self.return_random_features: model_output.append(gp_feature) return model_output class LaplaceRandomFeatureCovariance(tf_keras.layers.Layer): """Computes the Gaussian Process covariance using Laplace method. At training time, this layer updates the Gaussian process posterior using model features in minibatches. Attributes: momentum: (float) A discount factor used to compute the moving average for posterior precision matrix. Analogous to the momentum factor in batch normalization. If -1 then update covariance matrix using a naive sum without momentum, which is desirable if the goal is to compute the exact covariance matrix by passing through data once (say in the final epoch). ridge_penalty: (float) Initial Ridge penalty to weight covariance matrix. This value is used to stablize the eigenvalues of weight covariance estimate so that the matrix inverse can be computed for Cov = inv(t(X) * X + s * I). The ridge factor s cannot be too large since otherwise it will dominate the t(X) * X term and make covariance estimate not meaningful. likelihood: (str) The likelihood to use for computing Laplace approximation for the covariance matrix. Can be one of ('binary_logistic', 'poisson', 'gaussian'). """ def __init__(self, momentum=0.999, ridge_penalty=1., likelihood='gaussian', dtype=None, name='laplace_covariance'): if likelihood not in _SUPPORTED_LIKELIHOOD: raise ValueError( f'"likelihood" must be one of {_SUPPORTED_LIKELIHOOD}, got {likelihood}.' ) self.ridge_penalty = ridge_penalty self.momentum = momentum self.likelihood = likelihood super(LaplaceRandomFeatureCovariance, self).__init__(dtype=dtype, name=name) def compute_output_shape(self, input_shape): gp_feature_dim = input_shape[-1] return tf.TensorShape([gp_feature_dim, gp_feature_dim]) def build(self, input_shape): gp_feature_dim = input_shape[-1] # Convert gp_feature_dim to int value for TF1 compatibility. if isinstance(gp_feature_dim, tf.compat.v1.Dimension): gp_feature_dim = gp_feature_dim.value # Posterior precision matrix for the GP's random feature coefficients. self.initial_precision_matrix = ( self.ridge_penalty * tf.eye(gp_feature_dim, dtype=self.dtype)) self.precision_matrix = ( self.add_weight( name='gp_precision_matrix', shape=(gp_feature_dim, gp_feature_dim), dtype=self.dtype, initializer=tf_keras.initializers.Identity(self.ridge_penalty), trainable=False, aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)) self.built = True def make_precision_matrix_update_op(self, gp_feature, logits, precision_matrix): """Defines update op for the precision matrix of feature weights.""" if self.likelihood != 'gaussian': if logits is None: raise ValueError( f'"logits" cannot be None when likelihood={self.likelihood}') if logits.shape[-1] != 1: raise ValueError( f'likelihood={self.likelihood} only support univariate logits.' f'Got logits dimension: {logits.shape[-1]}') batch_size = tf.shape(gp_feature)[0] batch_size = tf.cast(batch_size, dtype=gp_feature.dtype) # Computes batch-specific normalized precision matrix. if self.likelihood == 'binary_logistic': prob = tf.sigmoid(logits) prob_multiplier = prob * (1. - prob) elif self.likelihood == 'poisson': prob_multiplier = tf.exp(logits) else: prob_multiplier = 1. gp_feature_adjusted = tf.sqrt(prob_multiplier) * gp_feature precision_matrix_minibatch = tf.matmul( gp_feature_adjusted, gp_feature_adjusted, transpose_a=True) # Updates the population-wise precision matrix. if self.momentum > 0: # Use moving-average updates to accumulate batch-specific precision # matrices. precision_matrix_minibatch = precision_matrix_minibatch / batch_size precision_matrix_new = ( self.momentum * precision_matrix + (1. - self.momentum) * precision_matrix_minibatch) else: # Compute exact population-wise covariance without momentum. # If use this option, make sure to pass through data only once. precision_matrix_new = precision_matrix + precision_matrix_minibatch # Returns the update op. return precision_matrix.assign(precision_matrix_new) def reset_precision_matrix(self): """Resets precision matrix to its initial value. This function is useful for reseting the model's covariance matrix at the beginning of a new epoch. """ precision_matrix_reset_op = self.precision_matrix.assign( self.initial_precision_matrix) self.add_update(precision_matrix_reset_op) def compute_predictive_covariance(self, gp_feature): """Computes posterior predictive variance. Approximates the Gaussian process posterior using random features. Given training random feature Phi_tr (num_train, num_hidden) and testing random feature Phi_ts (batch_size, num_hidden). The predictive covariance matrix is computed as (assuming Gaussian likelihood): s * Phi_ts @ inv(t(Phi_tr) * Phi_tr + s * I) @ t(Phi_ts), where s is the ridge factor to be used for stablizing the inverse, and I is the identity matrix with shape (num_hidden, num_hidden). Args: gp_feature: (tf.Tensor) The random feature of testing data to be used for computing the covariance matrix. Shape (batch_size, gp_hidden_size). Returns: (tf.Tensor) Predictive covariance matrix, shape (batch_size, batch_size). """ # Computes the covariance matrix of the feature coefficient. feature_cov_matrix = tf.linalg.inv(self.precision_matrix) # Computes the covariance matrix of the gp prediction. cov_feature_product = tf.matmul( feature_cov_matrix, gp_feature, transpose_b=True) * self.ridge_penalty gp_cov_matrix = tf.matmul(gp_feature, cov_feature_product) return gp_cov_matrix def _get_training_value(self, training=None): if training is None: training = tf_keras.backend.learning_phase() if isinstance(training, int): training = bool(training) return training def call(self, inputs, logits=None, training=None): """Minibatch updates the GP's posterior precision matrix estimate. Args: inputs: (tf.Tensor) GP random features, shape (batch_size, gp_hidden_size). logits: (tf.Tensor) Pre-activation output from the model. Needed for Laplace approximation under a non-Gaussian likelihood. training: (tf.bool) whether or not the layer is in training mode. If in training mode, the gp_weight covariance is updated using gp_feature. Returns: gp_stddev (tf.Tensor): GP posterior predictive variance, shape (batch_size, batch_size). """ batch_size = tf.shape(inputs)[0] training = self._get_training_value(training) if training: # Define and register the update op for feature precision matrix. precision_matrix_update_op = self.make_precision_matrix_update_op( gp_feature=inputs, logits=logits, precision_matrix=self.precision_matrix) self.add_update(precision_matrix_update_op) # Return null estimate during training. return tf.eye(batch_size, dtype=self.dtype) else: # Return covariance estimate during inference. return self.compute_predictive_covariance(gp_feature=inputs) def mean_field_logits(logits, covariance_matrix=None, mean_field_factor=1.): """Adjust the model logits so its softmax approximates the posterior mean [1]. [1]: Zhiyun Lu, Eugene Ie, Fei Sha. Uncertainty Estimation with Infinitesimal Jackknife. _arXiv preprint arXiv:2006.07584_, 2020. https://arxiv.org/abs/2006.07584 Arguments: logits: A float tensor of shape (batch_size, num_classes). covariance_matrix: The covariance matrix of shape (batch_size, batch_size). If None then it assumes the covariance_matrix is an identity matrix. mean_field_factor: The scale factor for mean-field approximation, used to adjust the influence of posterior variance in posterior mean approximation. If covariance_matrix=None then it is used as the temperature parameter for temperature scaling. Returns: Tensor of adjusted logits, shape (batch_size, num_classes). """ if mean_field_factor is None or mean_field_factor < 0: return logits # Compute standard deviation. if covariance_matrix is None: variances = 1. else: variances = tf.linalg.diag_part(covariance_matrix) # Compute scaling coefficient for mean-field approximation. logits_scale = tf.sqrt(1. + variances * mean_field_factor) if len(logits.shape) > 1: # Cast logits_scale to compatible dimension. logits_scale = tf.expand_dims(logits_scale, axis=-1) return logits / logits_scale