Add BetaPrior for correlation parameter priors (#3266)

Summary:
Pull Request resolved: #3266

X-link: facebook/Ax#5152

Adds `BetaPrior` — a Beta distribution prior supported on [0, 1], useful for correlation parameters in multi-task GPs. Registers it in Ax's FB serialization registries.

Reviewed By: sdaulton

Differential Revision: D99841563

fbshipit-source-id: 4064e1d0a303aa554086c1a633072b788212a97a

Author: Carl Hvarfner

botorch/models/utils/priors.py

@@ -0,0 +1,76 @@
+#!/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 __future__ import annotations
+
+from typing import Callable, Optional
+
+import torch
+from gpytorch.priors import Prior
+from gpytorch.priors.utils import BUFFERED_PREFIX
+from torch import Tensor
+from torch.distributions import Beta
+from torch.nn import Module as TModule
+
+
+class BetaPrior(Prior, Beta):
+    """Beta Prior parameterized by concentration1 (alpha) and concentration0 (beta).
+
+    pdf(x) = x^(alpha - 1) * (1 - x)^(beta - 1) / B(alpha, beta)
+
+    where alpha > 0 and beta > 0 are the concentration parameters.
+    Supported on [0, 1], useful as a prior on correlation parameters.
+    """
+
+    # Beta.concentration1/concentration0 are @property descriptors (they
+    # delegate to an internal Dirichlet), so _bufferize_attributes cannot
+    # delattr them.  We store separate buffers with the BUFFERED_PREFIX and
+    # sync them back after load_state_dict.
+    _PARAM_NAMES = ("concentration1", "concentration0")
+
+    def __init__(
+        self,
+        concentration1: float,
+        concentration0: float,
+        validate_args: bool = False,
+        transform: Optional[Callable[[Tensor], Tensor]] = None,
+    ) -> None:
+        """Initialize BetaPrior.
+
+        Args:
+            concentration1: Alpha (first concentration) parameter.
+            concentration0: Beta (second concentration) parameter.
+            validate_args: Whether to validate input arguments.
+            transform: Optional transform to apply before computing log_prob.
+        """
+        TModule.__init__(self)
+        Beta.__init__(
+            self,
+            concentration1=concentration1,
+            concentration0=concentration0,
+            validate_args=validate_args,
+        )
+        for attr in self._PARAM_NAMES:
+            self.register_buffer(
+                f"{BUFFERED_PREFIX}{attr}", getattr(self, attr).clone()
+            )
+        self._transform = transform
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        # Sync buffered values back into the underlying Dirichlet distribution.
+        c1 = getattr(self, f"{BUFFERED_PREFIX}concentration1")
+        c0 = getattr(self, f"{BUFFERED_PREFIX}concentration0")
+        self._dirichlet.concentration = torch.stack([c1, c0], dim=-1)
+
+    def expand(self, batch_shape: torch.Size) -> "BetaPrior":
+        batch_shape = torch.Size(batch_shape)
+        return BetaPrior(
+            self.concentration1.expand(batch_shape),
+            self.concentration0.expand(batch_shape),
+            validate_args=self._validate_args,
+            transform=self._transform,
+        )

sphinx/source/models.rst

@@ -209,6 +209,11 @@ Inducing Point Allocators
     :members:
     :private-members: _pivoted_cholesky_init
 
+Priors
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+.. automodule:: botorch.models.utils.priors
+    :members:
+
 Other Utilties
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. automodule:: botorch.models.utils.assorted

test/models/utils/test_priors.py

@@ -0,0 +1,84 @@
+#!/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 unittest
+
+import torch
+from botorch.models.utils.priors import BetaPrior
+from gpytorch.priors.utils import BUFFERED_PREFIX
+from torch.distributions import Beta
+
+
+class TestBetaPrior(unittest.TestCase):
+    def test_init(self):
+        prior = BetaPrior(1.2, 0.9)
+        self.assertAlmostEqual(prior.concentration1.item(), 1.2)
+        self.assertAlmostEqual(prior.concentration0.item(), 0.9)
+        self.assertIsNone(prior._transform)
+
+    def test_init_with_transform(self):
+        prior = BetaPrior(2.0, 3.0, transform=torch.sigmoid)
+        self.assertIs(prior._transform, torch.sigmoid)
+
+    def test_log_prob_matches_torch(self):
+        prior = BetaPrior(1.2, 0.9)
+        ref = Beta(torch.tensor(1.2), torch.tensor(0.9))
+        x = torch.rand(5, 4)
+        self.assertTrue(torch.allclose(prior.log_prob(x), ref.log_prob(x)))
+
+    def test_log_prob_with_transform(self):
+        def transform(x):
+            return x.clamp(0.01, 0.99)
+
+        prior = BetaPrior(2.0, 2.0, transform=transform)
+        prior_no_transform = BetaPrior(2.0, 2.0)
+        x = torch.rand(3, 2)
+        lp = prior.log_prob(x)
+        self.assertEqual(lp.shape, torch.Size([3, 2]))
+        # Values at boundaries should differ due to clamping
+        x_boundary = torch.tensor([[0.001, 0.999]])
+        self.assertFalse(
+            torch.allclose(
+                prior.log_prob(x_boundary),
+                prior_no_transform.log_prob(x_boundary),
+            )
+        )
+
+    def test_log_prob_batch(self):
+        prior = BetaPrior(1.5, 2.5)
+        x = torch.rand(7, 3, 2)
+        lp = prior.log_prob(x)
+        self.assertEqual(lp.shape, torch.Size([7, 3, 2]))
+
+    def test_rsample(self):
+        prior = BetaPrior(1.2, 0.9)
+        samples = prior.rsample(torch.Size([10, 5]))
+        self.assertEqual(samples.shape, torch.Size([10, 5]))
+        self.assertTrue(torch.all(samples >= 0))
+        self.assertTrue(torch.all(samples <= 1))
+
+    def test_state_dict_roundtrip(self):
+        prior = BetaPrior(1.2, 0.9)
+        sd = prior.state_dict()
+        self.assertIn(f"{BUFFERED_PREFIX}concentration1", sd)
+        self.assertIn(f"{BUFFERED_PREFIX}concentration0", sd)
+
+        prior2 = BetaPrior(999.0, 999.0)
+        prior2.load_state_dict(sd)
+        self.assertAlmostEqual(prior2.concentration1.item(), 1.2)
+        self.assertAlmostEqual(prior2.concentration0.item(), 0.9)
+
+    def test_expand(self):
+        prior = BetaPrior(1.2, 0.9)
+        expanded = prior.expand(torch.Size([3, 2]))
+        self.assertEqual(expanded.concentration1.shape, torch.Size([3, 2]))
+        self.assertEqual(expanded.concentration0.shape, torch.Size([3, 2]))
+
+    def test_expand_preserves_transform(self):
+        prior = BetaPrior(1.2, 0.9, transform=torch.sigmoid)
+        expanded = prior.expand(torch.Size([3, 2]))
+        self.assertIs(expanded._transform, torch.sigmoid)
+        self.assertEqual(expanded._validate_args, prior._validate_args)