PositiveIndexKernel (#3047)

Summary:

PositiveIndexKernel - a MultiTaskGP kernel that enforces positive correlation. Should probably be upstreamed into GPyTorch at some point.

Also introduces priors on diagonal and off-diagonals separately, so that priors can be set on task correlation in a more intuititve fashion.

Differential Revision: D84878629

Author: Carl Hvarfner

botorch/models/kernels/positive_index.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional
+
+import torch
+from gpytorch.constraints import GreaterThan
+from gpytorch.kernels import IndexKernel
+from gpytorch.priors import Prior
+
+
+class PositiveIndexKernel(IndexKernel):
+    r"""
+    A kernel for discrete indices with strictly positive correlations.
+
+    Similar to IndexKernel but ensures all off-diagonal correlations are positive
+    by using a Cholesky-like parameterization with positive elements.
+
+    .. math::
+        k(i, j) = \frac{(LL^T)_{i,j}}{(LL^T)_{t,t}}
+
+    where L is a lower triangular matrix with positive elements and t is the
+    target_task_index.
+
+    Args:
+        num_tasks (int): Total number of indices.
+        rank (int): Rank of the covariance matrix parameterization.
+        task_prior (Prior, optional): Prior for the covariance matrix.
+        diag_prior (Prior, optional): Prior for the diagonal elements.
+        var_constraint (Interval, optional): Constraint for variance (not used, kept for API compatibility).
+        normalize_covar_matrix (bool): Whether to normalize the covariance matrix.
+        target_task_index (int): Index of the task whose diagonal element should be
+            normalized to 1. Defaults to 0 (first task).
+        unit_scale_for_target (bool): Whether to ensure the target task's has unit otuputscale.
+        **kwargs: Additional arguments passed to IndexKernel.
+    """
+
+    def __init__(
+        self,
+        num_tasks: int,
+        rank: Optional[int] = 1,
+        task_prior: Optional[Prior] = None,
+        diag_prior: Optional[Prior] = None,
+        normalize_covar_matrix: bool = False,
+        target_task_index: int = 0,
+        unit_scale_for_target: bool = True,
+        **kwargs,
+    ):
+        if rank > num_tasks:
+            raise RuntimeError(
+                "Cannot create a task covariance matrix larger than the number of tasks"
+            )
+        if not (0 <= target_task_index < num_tasks):
+            raise ValueError(
+                f"target_task_index must be between 0 and {num_tasks - 1}, "
+                f"got {target_task_index}"
+            )
+        super().__init__(
+            num_tasks=num_tasks,
+            rank=rank,
+            prior=task_prior,
+            var_constraint=None,
+            **kwargs,
+        )
+        self.normalize_covar_matrix = normalize_covar_matrix
+        self.num_tasks = num_tasks
+        self.target_task_index = target_task_index
+        self.register_parameter(
+            name="raw_covar_factor",
+            parameter=torch.nn.Parameter(
+                torch.rand(*self.batch_shape, num_tasks, rank)
+            ),
+        )
+        self.unit_scale_for_target = unit_scale_for_target
+        if task_prior is not None:
+            if not isinstance(task_prior, Prior):
+                raise TypeError(
+                    f"Expected gpytorch.priors.Prior but got "
+                    f"{type(task_prior).__name__}"
+                )
+            self.register_prior(
+                "IndexKernelPrior", task_prior, lambda m: m._lower_triangle
+            )
+        if diag_prior is not None:
+            self.register_prior("ScalePrior", diag_prior, lambda m: m._diagonal)
+
+        self.register_constraint("raw_covar_factor", GreaterThan(0.0))
+
+    def _covar_factor_params(self, m):
+        return m.covar_factor
+
+    def _covar_factor_closure(self, m, v):
+        m._set_covar_factor(v)
+
+    @property
+    def covar_factor(self):
+        return self.raw_covar_factor_constraint.transform(self.raw_covar_factor)
+
+    @covar_factor.setter
+    def covar_factor(self, value):
+        self._set_covar_factor(value)
+
+    def _set_covar_factor(self, value):
+        # This must be a tensor
+        self.initialize(
+            raw_covar_factor=self.raw_covar_factor_constraint.inverse_transform(value)
+        )
+
+    @property
+    def _lower_triangle(self):
+        lower_row, lower_col = torch.tril_indices(
+            self.num_tasks, self.num_tasks, offset=-1
+        )
+        covar = self.covar_matrix
+        norm_factor = covar.diagonal(dim1=-1, dim2=-2).sqrt()
+        corr = covar / (norm_factor.unsqueeze(-1) * norm_factor.unsqueeze(-2))
+        low_tri = corr[..., lower_row, lower_col]
+
+        return low_tri
+
+    @property
+    def _diagonal(self):
+        return torch.diagonal(self.covar_matrix, dim1=-2, dim2=-1)
+
+    def _eval_covar_matrix(self):
+        cf = self.covar_factor
+        covar = cf @ cf.transpose(-1, -2) + self.var * torch.eye(
+            self.num_tasks, dtype=cf.dtype, device=cf.device
+        )
+        # Normalize by the target task's diagonal element
+        if self.unit_scale_for_target:
+            norm_factor = covar[..., self.target_task_index, self.target_task_index]
+            covar = covar / norm_factor.unsqueeze(-1).unsqueeze(-1)
+        return covar
+
+    @property
+    def covar_matrix(self):
+        return self._eval_covar_matrix()

test/models/kernels/test_positive_index.py

@@ -0,0 +1,199 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from botorch.models.kernels.positive_index import PositiveIndexKernel
+from botorch.utils.testing import BotorchTestCase
+from gpytorch.priors import NormalPrior
+
+
+class TestPositiveIndexKernel(BotorchTestCase):
+    def test_positive_index_kernel(self):
+        """Comprehensive test for PositiveIndexKernel functionality."""
+
+        # Test initialization
+        with self.subTest("basic_initialization"):
+            num_tasks = 4
+            rank = 2
+            kernel = PositiveIndexKernel(num_tasks=num_tasks, rank=rank)
+
+            self.assertEqual(kernel.num_tasks, num_tasks)
+            self.assertEqual(kernel.raw_covar_factor.shape, (num_tasks, rank))
+            self.assertEqual(kernel.normalize_covar_matrix, False)
+
+        # Test initialization with batch shape
+        with self.subTest("initialization_with_batch_shape"):
+            num_tasks = 3
+            rank = 2
+            batch_shape = torch.Size([2])
+            kernel = PositiveIndexKernel(
+                num_tasks=num_tasks, rank=rank, batch_shape=batch_shape
+            )
+
+            self.assertEqual(kernel.raw_covar_factor.shape, (2, num_tasks, rank))
+
+        # Test rank validation
+        with self.subTest("rank_validation"):
+            num_tasks = 3
+            rank = 5
+            with self.assertRaises(RuntimeError):
+                PositiveIndexKernel(num_tasks=num_tasks, rank=rank)
+
+        # Test target_task_index validation
+        with self.subTest("target_task_index_validation"):
+            num_tasks = 4
+            # Test invalid negative index
+            with self.assertRaises(ValueError):
+                PositiveIndexKernel(num_tasks=num_tasks, rank=2, target_task_index=-1)
+            # Test invalid index >= num_tasks
+            with self.assertRaises(ValueError):
+                PositiveIndexKernel(num_tasks=num_tasks, rank=2, target_task_index=4)
+            # Test valid indices (should not raise)
+            PositiveIndexKernel(num_tasks=num_tasks, rank=2, target_task_index=0)
+            PositiveIndexKernel(num_tasks=num_tasks, rank=2, target_task_index=3)
+
+        # Test covar_factor constraint
+        with self.subTest("positive_correlations"):
+            kernel = PositiveIndexKernel(num_tasks=5, rank=3)
+            covar_factor = kernel.covar_factor
+
+            # All elements should be positive
+            self.assertTrue((covar_factor > 0).all())
+
+            self.assertTrue((kernel.covar_matrix >= 0).all())
+
+        # Test covariance matrix normalization (default target_task_index=0)
+        with self.subTest("covar_matrix_normalization_default"):
+            kernel = PositiveIndexKernel(num_tasks=4, rank=2)
+            covar = kernel.covar_matrix
+
+            # First diagonal element should be 1.0 (normalized by default)
+            self.assertAllClose(covar[0, 0], torch.tensor(1.0), atol=1e-4)
+
+        # Test covariance matrix normalization with custom target_task_index
+        with self.subTest("covar_matrix_normalization_custom_target"):
+            kernel = PositiveIndexKernel(num_tasks=4, rank=2, target_task_index=2)
+            covar = kernel.covar_matrix
+
+            # Third diagonal element should be 1.0 (target_task_index=2)
+            self.assertAllClose(covar[2, 2], torch.tensor(1.0), atol=1e-4)
+
+            # Other diagonal elements should not be 1.0
+            self.assertNotEqual(covar[0, 0].item(), 1.0)
+
+        # Test forward pass shape
+        with self.subTest("forward"):
+            num_tasks = 4
+            kernel = PositiveIndexKernel(num_tasks=num_tasks, rank=2)
+            kernel.eval()
+
+            # Create index inputs
+            i1 = torch.tensor([[0, 1], [2, 3]], dtype=torch.long)
+            i2 = torch.tensor([[1, 2]], dtype=torch.long)
+
+            result = kernel(i1, i2)
+
+            # Expected shape: (2, 2, 1, 1)
+            self.assertEqual(result.shape, torch.Size([2, 1]))
+            num_tasks = 3
+            kernel = PositiveIndexKernel(num_tasks=num_tasks, rank=1)
+            kernel.eval()
+
+            # Set specific covar_factor values for predictable output
+            kernel.initialize(raw_covar_factor=torch.ones(num_tasks, 1))
+
+            i1 = torch.tensor([[0]], dtype=torch.long)
+            i2 = torch.tensor([[1]], dtype=torch.long)
+
+            result = kernel(i1, i2).to_dense()
+            covar_matrix = kernel.covar_matrix
+            expected = covar_matrix[0, 1]
+
+            self.assertAllClose(result.squeeze(), expected)
+
+        # Test with priors
+        with self.subTest("with_priors"):
+            num_tasks = 4
+            task_prior = NormalPrior(0, 1)
+            diag_prior = NormalPrior(1, 0.1)
+
+            kernel = PositiveIndexKernel(
+                num_tasks=num_tasks,
+                rank=2,
+                task_prior=task_prior,
+                diag_prior=diag_prior,
+                initialize_to_mode=False,
+            )
+
+            # Check that priors are registered
+            prior_names = [p[0] for p in kernel.named_priors()]
+            self.assertIn("IndexKernelPrior", prior_names)
+            self.assertIn("ScalePrior", prior_names)
+
+        # Test batch forward
+        with self.subTest("batch_forward"):
+            num_tasks = 3
+            batch_shape = torch.Size([2])
+            kernel = PositiveIndexKernel(
+                num_tasks=num_tasks, rank=2, batch_shape=batch_shape
+            )
+            kernel.eval()
+
+            i1 = torch.tensor([[[0], [1]]], dtype=torch.long)
+            i2 = torch.tensor([[[1], [2]]], dtype=torch.long)
+
+            result = kernel(i1, i2)
+
+            # Check that batch dimensions are preserved
+            self.assertEqual(result.shape[0], 2)
+
+        # Test diagonal property (default target_task_index=0)
+        with self.subTest("diagonal"):
+            kernel = PositiveIndexKernel(num_tasks=4, rank=2)
+            diag = kernel._diagonal
+
+            self.assertEqual(diag.shape, torch.Size([4]))
+            # First diagonal element should be 1.0 (default target_task_index=0)
+            self.assertAllClose(diag[0], torch.tensor(1.0), atol=1e-4)
+
+            # Test diagonal property with custom target_task_index
+            kernel = PositiveIndexKernel(num_tasks=4, rank=2, target_task_index=1)
+            diag = kernel._diagonal
+
+            self.assertEqual(diag.shape, torch.Size([4]))
+            # Second diagonal element should be 1.0 (target_task_index=1)
+            self.assertAllClose(diag[1], torch.tensor(1.0), atol=1e-4)
+
+        # Test lower triangle property
+        with self.subTest("lower_triangle"):
+            num_tasks = 5
+            kernel = PositiveIndexKernel(num_tasks=num_tasks, rank=2)
+            lower_tri = kernel._lower_triangle
+
+            # Number of lower triangular elements (excluding diagonal)
+            expected_size = num_tasks * (num_tasks - 1) // 2
+            self.assertEqual(lower_tri.shape[-1], expected_size)
+            self.assertTrue((lower_tri >= 0).all())
+
+        # Test invalid prior type
+        with self.subTest("invalid_prior_type"):
+            with self.assertRaises(TypeError):
+                PositiveIndexKernel(num_tasks=4, rank=2, task_prior="not_a_prior")
+
+        # Test covariance matrix properties
+        with self.subTest("covar_matrix"):
+            kernel = PositiveIndexKernel(num_tasks=5, rank=4)
+            covar = kernel.covar_matrix
+
+            # Should be square
+            self.assertEqual(covar.shape[-2], covar.shape[-1])
+
+            # Should be positive definite (all eigenvalues > 0)
+            eigvals = torch.linalg.eigvalsh(covar)
+            self.assertTrue((eigvals > 0).all())
+
+            # Should be symmetric
+            self.assertAllClose(covar, covar.T, atol=1e-5)