GPA layer and NLL loss functiona added by hasan2m · Pull Request #72 · JeffersonLab/jlab_datascience_core

hasan2m · 2025-03-28T04:39:02Z

Unit test passed for both GPA layer and loss function.

GPA layer:

Loss function:

sgoldenCS

I generally think things look fine although I haven't run the unittests myself (maybe someone else could go through that process?). I'm glad you're using pytest and pytest fixtures though! It looks good!

sgoldenCS · 2025-03-28T13:46:04Z

+            initializer=tf.constant_initializer(self.initial_noise_scale),
+            trainable=self.train_noise_scale,
+            dtype=tf.float32,
+            constraint=ClipByValue(1e-6, 1e6),


I like the use of a function for the constraint, but would it be possible to continue using the noise_bounds parameter here? Otherwise it should be removed from the class definition as I don't believe it's used elsewhere.

Removed noise_bound from the class definition.

sgoldenCS · 2025-03-28T13:47:05Z

+        # self.noise_scale = tf.Variable(
+        #     self.initial_noise_scale,
+        #     dtype=tf.float32,
+        #     trainable=self.train_noise_scale,
+        #     constraint=lambda z: tf.clip_by_value(z, self.noise_bounds[0], self.noise_bounds[1]),
+        #     name='noise_scale'
+        # )


This looks to be essentially the same as the new version (other than the noise_bounds thing), so feel free to remove the commented section to clean things up.

sgoldenCS · 2025-03-28T13:49:31Z

+    # def call(self, inputs, training=None, return_features=False):
+    #     if training is None:
+    #         training = tf.keras.backend.learning_phase()
+
+    #     batch_size = tf.cast(tf.shape(inputs)[0], tf.float32)
+    #     x = tf.convert_to_tensor(inputs, dtype=self.dtype)
+    #     x = tf.cast(x, tf.float32)
+
+    #     x = self.length_scale * x
+    #     x = self.rff_map(x)
+    #     x1 = tf.math.cos(x) 
+    #     x2 = tf.math.sin(x) 
+
+    #     ffs = layers.concatenate([x1, x2])
+
+    #     if self.scale_features:
+    #         ffs = tf.math.sqrt(2.0 / self.n_fourier_features) * ffs
+
+    #     ffs = tf.math.sqrt(self.constant_scale) * ffs
+    #     output = self.rff_output(ffs)
+
+    #     if training:
+    #         if self.momentum > 0:
+    #             update_prior_op = (
+    #                 self.momentum * self.prior + (1 - self.momentum) * (tf.transpose(ffs) @ ffs / batch_size)
+    #             )
+    #         else:
+    #             update_prior_op = self.prior + tf.transpose(ffs) @ ffs
+    #         self.prior.assign(update_prior_op)  # Direct assignment
+
+    #         variances = self.calc_variance(ffs)
+    #     else:
+    #         if not self.do_custom_cov_update:
+    #             self.update_cov(self.prior)
+
+    #         variances = self.calc_variance(ffs)
+
+    #     stddevs = tf.math.sqrt(variances)
+    #     out = [output, stddevs[:, None]]
+    #     if return_features:
+    #         out.append(ffs)
+
+    #     return out


I don't see any changes here other than the default for the training flag, so I think this commented section could be removed too.

clean up code.

schr476 · 2025-03-28T13:58:46Z

GaussianProcessLayer --> We should rename since this is RFF approximation and not GP

Name changed to GaussianProcessApproxiamtionLayer.

remove this from utils.

schr476 · 2025-04-01T18:09:23Z

please change file name --> gpa_layer.py

schr476 · 2025-04-01T18:11:54Z

+        tf.Tensor: The computed NLL loss (scalar).
+    """
+
+    sigma_star = tf.square(std) + noise_scale + 1e-5  # Adding a small constant for numerical stability


is the noise scale additive?

Removing the noise_scale from here because we are accounting the noise scale during the variance calculation within the gpa_layer class.

schr476 · 2025-04-01T18:14:24Z

+import numpy as np
+import tensorflow as tf
+from tensorflow.keras import Model, Input
+from jlab_datascience_toolkit.utils.keras_layers.GP_layer import GaussianProcessLayer


update with request change above.
Please don't call it a GaussianProcessLayer.

schr476 · 2025-04-01T18:14:46Z

+
+
+def test_forward_pass(random_input):
+    layer = GaussianProcessLayer()


Please don't call it GaussianProcessLayer

hasan2m · 2025-04-03T03:52:17Z

Addressed the comments and reran the unit test.

Kishanrajput · 2025-04-17T20:12:01Z

@hasan2m - could you please move the layer and loss out from utils by following below directory structure
Toolkit

keras (dir)
- layers
  - layer_v0.py (implementation)
- losses (dir)
  - loss_v0.py (implementation)

schr476 · 2025-04-17T20:10:36Z

remove this from utils.

schr476 · 2025-04-17T20:11:31Z

+    # def call(self, inputs, training=None, return_features=False):
+    #     if training is None:
+    #         training = tf.keras.backend.learning_phase()
+
+    #     batch_size = tf.cast(tf.shape(inputs)[0], tf.float32)
+    #     x = tf.convert_to_tensor(inputs, dtype=self.dtype)
+    #     x = tf.cast(x, tf.float32)
+
+    #     x = self.length_scale * x
+    #     x = self.rff_map(x)
+    #     x1 = tf.math.cos(x) 
+    #     x2 = tf.math.sin(x) 
+
+    #     ffs = layers.concatenate([x1, x2])
+
+    #     if self.scale_features:
+    #         ffs = tf.math.sqrt(2.0 / self.n_fourier_features) * ffs
+
+    #     ffs = tf.math.sqrt(self.constant_scale) * ffs
+    #     output = self.rff_output(ffs)
+
+    #     if training:
+    #         if self.momentum > 0:
+    #             update_prior_op = (
+    #                 self.momentum * self.prior + (1 - self.momentum) * (tf.transpose(ffs) @ ffs / batch_size)
+    #             )
+    #         else:
+    #             update_prior_op = self.prior + tf.transpose(ffs) @ ffs
+    #         self.prior.assign(update_prior_op)  # Direct assignment
+
+    #         variances = self.calc_variance(ffs)
+    #     else:
+    #         if not self.do_custom_cov_update:
+    #             self.update_cov(self.prior)
+
+    #         variances = self.calc_variance(ffs)
+
+    #     stddevs = tf.math.sqrt(variances)
+    #     out = [output, stddevs[:, None]]
+    #     if return_features:
+    #         out.append(ffs)
+
+    #     return out


clean up code.

schr476 · 2025-04-17T20:40:46Z

+
+        if training:
+            update_prior_op = (
+                self.momentum * self.prior + (1 - self.momentum) * (tf.transpose(ffs) @ ffs / batch_size)


check of the size is correct: (batch,batch)

schr476 · 2025-04-17T20:46:49Z

+        self.eigvals.assign(eigvals)
+        self.eigvecs.assign(eigvecs)
+
+    def calc_variance(self, ffs):


need to understand this better. this should be K_xx^{-1}

Inverse of K_xx is determined using eigenvalue and eigenvectors.

schr476 · 2025-04-17T21:01:29Z

+            name='rff_map'
+        )
+
+       self.rff_output = layers.Dense(self.n_out, use_bias=False, name='GP_mean_pred')


What is this here?

Output layer that provides the mean prediction using the fourier features created by rff_map.

schr476 · 2025-04-17T21:13:05Z

+        tf.Tensor: The computed NLL loss (scalar).
+    """
+
+    sigma_star = tf.square(std) + 1e-5  # Adding a small constant for numerical stability


sigma_star should be sigma2_star

schr476 · 2025-05-07T17:00:35Z

@hasan2m what is the status of the requested changes?

sgoldenCS · 2025-05-07T17:26:59Z

+    def update_cov(self, prior):
+        eigvals, eigvecs = tf.linalg.eigh(prior)  #EigenDecomposition
+        eigvals = tf.where(eigvals > 0, eigvals, tf.zeros_like(eigvals))
+
+        self.eigvals.assign(eigvals)
+        self.eigvecs.assign(eigvecs)


I wonder if this should be an internal version (_update_cov(self, prior) or something similar) that can be called with any prior, and implement the update_cov function without the extra input parameter like this:

def update_cov(self): self._update_cov(self.prior)

That could make it easier for a user to update the variance calculations outside of the model. Going along with this, it might also make sense to just remove the update_cov() call on line 146, and just rely on the user to properly update the prior when desired. At the moment, if do_custom_cov_update == False, the code will be slow for every non-training/validation call, which probably isn't desirable.

Kishanrajput

When a model that uses this layer is saved and loaded back - it does not save/load prior matrix which makes the saved models not useful. Please add prior saving and loading.

Kishanrajput · 2025-05-15T19:53:27Z

@hasan2m - please have a look at the comments and address/comment on them.

Kishanrajput · 2025-05-20T16:49:24Z

Print statement for the RFF shape should be removed

GP layer and NLL loss functiona added

359906c

hasan2m linked an issue Mar 28, 2025 that may be closed by this pull request

Add Gaussian Process Approximation (GPA) layer #71

Open

4 tasks

hasan2m requested review from Kishanrajput, schr476 and sgoldenCS March 28, 2025 04:39

hasan2m added the enhancement New feature or request label Mar 28, 2025

hasan2m self-assigned this Mar 28, 2025

sgoldenCS reviewed Mar 28, 2025

View reviewed changes

schr476 reviewed Apr 1, 2025

View reviewed changes

schr476 requested changes Apr 1, 2025

View reviewed changes

addressed the comments on PR

5f5a397

schr476 requested changes Apr 17, 2025

View reviewed changes

New directory

7d594fe

Kishanrajput requested review from schr476 and sgoldenCS April 29, 2025 17:42

Minor change

ec5edac

sgoldenCS reviewed May 7, 2025

View reviewed changes

Kishanrajput requested changes May 15, 2025

View reviewed changes



		def test_forward_pass(random_input):
		layer = GaussianProcessLayer()

Conversation

hasan2m commented Mar 28, 2025

Uh oh!

sgoldenCS left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hasan2m commented Apr 3, 2025

Uh oh!

Kishanrajput commented Apr 17, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

schr476 commented May 7, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Kishanrajput left a comment

Choose a reason for hiding this comment

Uh oh!

Kishanrajput commented May 15, 2025

Uh oh!

Kishanrajput commented May 20, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants