Implementation of the NTK in Colibri

colibri/app.py

@@ -32,6 +32,10 @@
     "colibri.param_initialisation",
     "colibri.export_results",
     "colibri.closure_test",
+    "colibri.ntk.eigenvalues",
+    "colibri.ntk.eigenvector",
+    "colibri.ntk.plotntk",
+    "colibri.ntk.ntk",
     "reportengine.report",
 ]
 

colibri/config.py

@@ -10,6 +10,7 @@
 import logging
 import os
 import shutil
+from pathlib import Path
 
 import jax
 import jax.numpy as jnp
@@ -18,10 +19,11 @@
 from colibri.constants import FLAVOUR_TO_ID_MAPPING
 from colibri.core import IntegrabilitySettings, PriorSettings
 from mpi4py import MPI
-from reportengine.configparser import ConfigError, explicit_node
+from reportengine.configparser import ConfigError, explicit_node, element_of
 from validphys import covmats
-from validphys.config import Config, Environment
+from validphys.config import Config, Environment, _id_with_label
 from validphys.fkparser import load_fktable
+from validphys.loader import LoadFailedError
 
 comm = MPI.COMM_WORLD
 rank = comm.Get_rank()
@@ -622,3 +624,20 @@ def produce_pdf_model(self):
         Returns None as the pdf_model is not used in the colibri module.
         """
         return None
+
+    def produce_replicas_path(self, fit):
+        """
+        Produces the replicas folder where the fit replicas are stored.
+        """
+        replicas_path = fit.path / "fit_replicas"
+        replicas_path.mkdir(parents=True, exist_ok=True)
+        return replicas_path
+
+    @element_of("fits")
+    @_id_with_label
+    def parse_fit(self, fit: str):
+        """A fit in the results folder, containing at least a valid filter result."""
+        try:
+            return self.loader.check_fit(fit)
+        except LoadFailedError as e:
+            raise ConfigError(str(e), fit, self.loader.available_fits)

colibri/core.py

@@ -167,12 +167,15 @@ class MonteCarloFit:
         Array containing the validation loss.
     optimized_parameters: jnp.array
         Array containing the optimized parameters.
+    parameters_by_epoch: jnp.array
+        Array containing the parameters of the model recorded during training.
     """
 
     monte_carlo_specs: dict
     training_loss: jnp.array
     validation_loss: jnp.array
     optimized_parameters: jnp.array
+    parameters_by_epoch: jnp.array
 
 
 @dataclass(frozen=True)
@@ -189,12 +192,15 @@ class GradientDescentResult:
         Recorded (epoch) validation losses (sampled according to record_every).
     specs: dict
         Dictionary of settings used for the run (epochs, batch size, etc.).
+    parameters_by_epoch: jnp.array
+        Array containing the parameters of the model recorded during training.
     """
 
     optimized_parameters: Any
     training_loss: jnp.array
     validation_loss: jnp.array
     specs: Dict[str, Any]
+    parameters_by_epoch: jnp.array
 
 
 @dataclass(frozen=True)

colibri/doc/sphinx/source/theory/index.rst

@@ -18,3 +18,5 @@ This section discusses some relevant theoretical background to Colibri.
    ./prior_distributions.rst
 
    ./inference_methods.rst
+
+   ./ntk_theory.rst

colibri/doc/sphinx/source/theory/ntk_theory.rst

@@ -0,0 +1,6 @@
+.. _ntk_theory:
+
+===========================
+Neural Tangent Kernel (NTK)
+===========================
+

colibri/doc/sphinx/source/tutorials/index.rst

@@ -17,3 +17,5 @@ such as implementing a custom PDF model or using it to fit data.
    closure_tests/index
 
    scripts/index
+
+   ntk/computing_ntk

colibri/doc/sphinx/source/tutorials/ntk/computing_ntk.rst

@@ -0,0 +1,79 @@
+.. _computing_ntk:
+
+=====================================
+Computing Neural Tangent Kernel (NTK)
+=====================================
+
+This section describes how to compute and analyse this Neural Tangent Kernel (NTK)
+during the training of a given fit. For information on what the NTK is, please
+refer to the :ref:`NTK theory section <ntk_theory>`.
+
+The NTK can be computed on any fit where the parameter values have been saved during
+training. To produce these saved parameter values during a fit set the
+`record_parameters` argument to `True` in the fit runcard.
+
+.. code-block:: bash
+
+    record_parameters: True
+    record_every: 100
+
+This above example will save the parameter values every 100 epochs.
+
+Once the parameter values have been saved, the NTK can be computed using the `compute_ntk`
+action. An example runcard to compute the NTK is shown below.
+
+.. code-block:: bash
+
+    meta:
+        title: Example NTK determination for a Les Houches fit         
+        author: Colibri User
+        keywords: [example, PDF plots, NTK]  
+
+    pdf: {id: "test_ntk_fit"} # generated by the colibri fit
+
+    ntk_plots_settings:
+        ntk_plots: True
+        n_top_eigenvalues: 10           # Number of top eigenvalues for which evolution is plotted
+        y_scale: log                    # Sets the scale for eigenvalue evolution plot.
+        plot_n_epochs: 5                # The number of epochs for which the eigenvector plots will be plotted.
+                                        # These will be equally spaced among all epochs. Default 6.
+        plot_n_eigenvectors: [1, 2]     # Eigenvector plots will be produced for these. Default 1.
+
+    actions_:
+    - compute_ntk
+
+This runcard can be run with the command
+``colibri_model_exe compute_ntk.yaml -r replica_n``, where the ``colibri_model_exe``
+is the Colibri executable, and ``n`` is the replica number for which the NTK
+should be computed. This will produce a folder called ``compute_ntk``.  The NTK values
+are stored in npz file format in ``compute_ntk/ntk_replicas/replica_n``.
+
+``ntk_plots_settings``
+^^^^^^^^^^^^^^^^^^^^^^
+* ``ntk_plots``: 
+    Whether analytic plots are produced with the ``compute_ntk`` command. For details on
+    the plots produced, see the :ref:`analytic NTK plots section below <analytic_ntk_plots>`.
+* ``n_top_eigenvalues``:
+    This is the number of top eigenvalues that will be plotted when ``ntk_plots`` is
+    ``True``. If a number is not specified, a default of 1 eigenvalue will be plotted.
+* ``x_scale`` / ``y_scale``:
+    If set to ``log``, the eigenvalue evolution plot will be produced with a logarithmic
+    scale on the respective axis. Otherwise a linear scale will be used.
+* ``plot_n_epochs``:
+    The number of epochs for which eigenvector plots will be produced.
+    These will be equally spaced among all epochs. Default 6.
+* ``plot_n_eigenvectors``:
+    The eigenvectors for which eigenvector plots will be produced. Default 1.
+
+.. _analytic_ntk_plots:
+
+Analytic NTK plots
+==================
+If the flag ``ntk_plots`` is set to ``True``, the following analytic plots will be
+produced and stored in .pdf format in ``compute_ntk/ntk_plots/``.
+
+* **NTK eigenvalue evolution:** Value of the NTK eigenvalues for increasing number of recorded epochs.
+
+* **NTK eigenvector evolution:** Eigenvector evolution for a given number of recorded epochs, for all flavours. 
+
+* **Eigenvector heatmaps:** Density of NTK eigenvectors across PDF flavours and x-regions.

colibri/examples/les_houches_example/runcards/lh_fit_closure_test.yaml

@@ -98,4 +98,4 @@ ultranest_settings:
 
 
 actions_:
-- run_ultranest_fit                      # Choose from ultranest_fit, monte_carlo_fit, analytic_fit
+- run_ultranest_fit                      # Choose from ultranest_fit, monte_carlo_fit, analytic_fit

colibri/gradient_descent.py

@@ -35,6 +35,7 @@ def run_gradient_descent(
     max_epochs: int,
     data_batch: colibri.DataBatches = None,
     record_every: int = 50,
+    record_parameters: bool = False,
 ) -> GradientDescentResult:
     """Generic gradient descent loop.
 
@@ -64,7 +65,11 @@ def run_gradient_descent(
         Defaults to None, in which case the loss is assumed to not have been batched.
 
     record_every : int, default 50
-        Record losses every this many epochs.
+        Record losses every this many epochs. If `record_parameters` is True,
+        parameters of the model are also recorded.
+
+    record_parameters : bool, default False
+        Whether to record parameters every `record_every` epochs.
     """
 
     params = initial_parameters
@@ -79,6 +84,7 @@ def _step(p, ostate, batch_idx):
 
     train_losses = []
     val_losses = []
+    parameters_by_epoch = [] if record_parameters else None
 
     if data_batch is None:
         # we simulate a fake batch iterator that just yields a dummy batch index
@@ -102,6 +108,7 @@ def batch_gen():
             batch_idx = next(batches_iter)
             params, opt_state, batch_loss = _step(params, opt_state, batch_idx)
             epoch_train_loss += batch_loss
+
         epoch_val_loss = validation_loss_fn(params)
 
         early_stopper = early_stopper.update(epoch_val_loss)
@@ -114,6 +121,11 @@ def batch_gen():
             train_losses.append(epoch_train_loss)
             val_losses.append(epoch_val_loss)
 
+        if record_parameters:
+            if epoch % record_every == 0:
+                log.info(f"Recording parameters at epoch {epoch}")
+                parameters_by_epoch.append(params)
+
         if early_stopper.should_stop:
             log.info(f"Early stopping at epoch {epoch}")
             break
@@ -127,4 +139,5 @@ def batch_gen():
             "batch_size": batch_size,
             "record_every": record_every,
         },
+        parameters_by_epoch=jnp.array(parameters_by_epoch),
     )

colibri/monte_carlo_fit.py

@@ -30,6 +30,8 @@ def monte_carlo_fit(
     max_epochs,
     batch_size=None,
     batch_seed=1,
+    record_parameters=False,
+    record_every=50,
 ):
     """
     This function performs a Monte Carlo fit.
@@ -61,14 +63,20 @@ def monte_carlo_fit(
     batch_seed: int, optional
         Seed used to construct the batches. Defaults to 1.
 
+    record_parameters: bool, default False
+        Whether to monitor the parameters during the Monte Carlo fit.
+
+    record_every: int, default 50
+        Frequency (in epochs) at which to record losses parameters of
+        model. If record_parameters is False, only losses are recorded.
+
     Returns
     -------
     MonteCarloFit: The result of the fit with following attributes:
         monte_carlo_specs: dict
         training_loss: jnp.array
         validation_loss: jnp.array
     """
-
     len_tr_idx, len_val_idx = len_trval_data
 
     @jax.jit
@@ -100,7 +108,8 @@ def loss_validation(parameters):
         early_stopper=early_stopper,
         max_epochs=max_epochs,
         data_batch=data_batch,
-        record_every=50,
+        record_every=record_every,
+        record_parameters=record_parameters,
     )
 
     t1 = time.time()
@@ -111,14 +120,18 @@ def loss_validation(parameters):
             "max_epochs": max_epochs,
             "batch_size": data_batch.batch_size,
             "batch_seed": batch_seed,
+            "record_every": record_every,
         },
         training_loss=gd_result.training_loss,
         validation_loss=gd_result.validation_loss,
         optimized_parameters=gd_result.optimized_parameters,
+        parameters_by_epoch=gd_result.parameters_by_epoch,
     )
 
 
-def run_monte_carlo_fit(monte_carlo_fit, pdf_model, output_path, replica_index, Q0):
+def run_monte_carlo_fit(
+    monte_carlo_fit, pdf_model, output_path, replica_index, Q0, record_parameters=False
+):
     """
     Runs the Monte Carlo fit and writes the output to the output directory.
 
@@ -178,3 +191,20 @@ def run_monte_carlo_fit(monte_carlo_fit, pdf_model, output_path, replica_index,
         index=False,
         float_format="%.5e",
     )
+
+    if record_parameters:
+        # Save the parameters by epoch if recorded
+        record_every = mc_fit.monte_carlo_specs["record_every"]
+        parameters_by_epoch = mc_fit.parameters_by_epoch
+        params_path = (
+            str(output_path) + f"/fit_replicas/replica_{replica_index}/parameters/"
+        )
+        if not os.path.exists(params_path):
+            os.makedirs(params_path, exist_ok=True)
+
+        for epoch_idx in range(parameters_by_epoch.shape[0]):
+            epoch = epoch_idx * record_every
+            jnp.savez(
+                params_path + f"params_{epoch}.npz",
+                params=parameters_by_epoch[epoch_idx],
+            )

colibri/ntk/__init__.py

@@ -0,0 +1,34 @@
+"""
+colibri.ntk
+
+NTK (Neural Tangent Kernel) analysis module for colibri.
+
+This module provides tools for computing and analyzing the Neural Tangent Kernel
+of PDF fits, including eigenvalue/eigenvector computation and plotting utilities.
+"""
+
+from colibri.ntk.ntkutils import NTKGrid, NTKStats
+from colibri.ntk.eigenvalues import (
+    EigenvalueGrid,
+    eigenvalue_grid,
+    eigenvalues_ensemble,
+)
+from colibri.ntk.eigenvector import (
+    EigenvectorGrid,
+    eigenvector_grid,
+    eigenvectors_ensemble_at_epoch,
+)
+
+__all__ = [
+    # Abstract interface
+    "NTKGrid",
+    "NTKStats",
+    # Eigenvalue classes and functions
+    "EigenvalueGrid",
+    "eigenvalue_grid",
+    "eigenvalues_ensemble",
+    # Eigenvector classes and functions
+    "EigenvectorGrid",
+    "eigenvector_grid",
+    "eigenvectors_ensemble_at_epoch",
+]