Refactor inducing point allocator classes to only need what they need (…

…facebookresearch#478)

Summary:
Pull Request resolved: facebookresearch#478

Inducing point allocator classes had extra methods and attributes that were not necessary or misnamed. These have been cleaned up in the BaseAllocator class and all of its children are similarly updated.

The dummy allocator is no longer needed, instead, the allocators can just make dummy points based on this dimensionality. This sets the last_allocator_used attribute to be None to signify no actual allocator was used in creating the last set of points.

This also requires all allocators to know its dimensionality at least as it is used by the dummy allocator when there's no inputs.

All models that use an allocator now initializes the allocators outside as there's no trivial default for allocators anymore.

Differential Revision: D67059839

aepsych/benchmark/example_problems.py

@@ -110,7 +110,7 @@ def __init__(
         self.m = GPClassificationModel(
             dim=6,
             inducing_size=inducing_size,
-            inducing_point_method=KMeansAllocator(bounds=self.bounds),
+            inducing_point_method=KMeansAllocator(dim=6),
         )
 
         self.m.fit(

aepsych/models/gp_classification.py

@@ -108,7 +108,6 @@ def __init__(
             inducing_size=self.inducing_size,
             covar_module=covar_module or default_covar,
         )
-        self.last_inducing_points_method = self.inducing_point_method.allocator_used
 
         self.inducing_points = inducing_points
         variational_distribution = CholeskyVariationalDistribution(
@@ -212,7 +211,6 @@ def _reset_variational_strategy(self) -> None:
                 covar_module=self.covar_module,
                 X=self.train_inputs[0],
             ).to(device)
-            self.last_inducing_points_method = self.inducing_point_method.allocator_used
             variational_distribution = CholeskyVariationalDistribution(
                 inducing_points.size(0), batch_shape=torch.Size([self._batch_size])
             ).to(device)
@@ -249,8 +247,7 @@ def fit(
             self._reset_hyperparameters()
 
         if not warmstart_induc or (
-            self.last_inducing_points_method == "DummyAllocator"
-            and self.inducing_point_method.__class__.__name__ != "DummyAllocator"
+            self.inducing_point_method.last_allocator_used is None
         ):
             self._reset_variational_strategy()
 

aepsych/models/inducing_points/__init__.py

@@ -9,15 +9,13 @@
 
 from ...config import Config
 from .auto import AutoAllocator
-from .base import DummyAllocator
 from .fixed import FixedAllocator
 from .greedy_variance_reduction import GreedyVarianceReduction
 from .kmeans import KMeansAllocator
 from .sobol import SobolAllocator
 
 __all__ = [
     "AutoAllocator",
-    "DummyAllocator",
     "FixedAllocator",
     "GreedyVarianceReduction",
     "KMeansAllocator",

aepsych/models/inducing_points/auto.py

@@ -1,123 +1,42 @@
-from typing import Any, Dict, Optional
+from typing import Optional
 
 import torch
-from aepsych.config import Config
-from aepsych.models.inducing_points.base import BaseAllocator, DummyAllocator
+from aepsych.models.inducing_points.base import BaseAllocator
 from aepsych.models.inducing_points.kmeans import KMeansAllocator
-from botorch.models.utils.inducing_point_allocators import InducingPointAllocator
 
 
 class AutoAllocator(BaseAllocator):
     """An inducing point allocator that dynamically chooses an allocation strategy
     based on the number of unique data points available."""
 
-    def __init__(
-        self,
-        bounds: Optional[torch.Tensor] = None,
-        fallback_allocator: InducingPointAllocator = KMeansAllocator(),
-    ) -> None:
-        """
-        Initialize the AutoAllocator with a fallback allocator.
-
-        Args:
-            fallback_allocator (InducingPointAllocator, optional): Allocator to use if there are
-                                                        more unique points than required.
-        """
-        super().__init__(bounds=bounds)
-        self.fallback_allocator = fallback_allocator
-        if bounds is not None:
-            self.bounds = bounds
-            self.dummy_allocator = DummyAllocator(bounds=bounds)
-
-    def _get_quality_function(self) -> None:
-        """AutoAllocator does not require a quality function, so this returns None."""
-        return None
-
     def allocate_inducing_points(
         self,
-        inputs: Optional[torch.Tensor],
+        inputs: Optional[torch.Tensor] = None,
         covar_module: Optional[torch.nn.Module] = None,
-        num_inducing: int = 10,
+        num_inducing: int = 100,
         input_batch_shape: torch.Size = torch.Size([]),
     ) -> torch.Tensor:
-        """
-        Allocate inducing points by either using the unique input data directly
-        or falling back to another allocation method if there are too many unique points.
+        """Generates `num_inducing` inducing points smartly based on the inputs.
+        Currently, this is just a wrapper for the KMeansAllocator
 
         Args:
             inputs (torch.Tensor): A tensor of shape (n, d) containing the input data.
             covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here.
-            num_inducing (int, optional): The number of inducing points to generate.
+            num_inducing (int, optional): The number of inducing points to generate. Defaults to 100.
             input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here.
 
         Returns:
-            torch.Tensor: A (num_inducing, d)-dimensional tensor of inducing points.
+            torch.Tensor: A (num_inducing, d)-dimensional tensor of inducing points selected via k-means++.
         """
-        # Ensure inputs are not None
-        if inputs is None and self.bounds is not None:
-            self.allocator_used = self.dummy_allocator.__class__.__name__
-            return self.dummy_allocator.allocate_inducing_points(
-                inputs=inputs,
-                covar_module=covar_module,
-                num_inducing=num_inducing,
-                input_batch_shape=input_batch_shape,
-            )
-        elif inputs is None and self.bounds is None:
-            raise ValueError(f"Either inputs or bounds must be provided.{self.bounds}")
-
-        assert (
-            inputs is not None
-        ), "inputs should not be None here"  # to make mypy happy
-
-        unique_inputs = torch.unique(inputs, dim=0)
-
-        # If there are fewer unique points than required, return unique inputs directly
-        if unique_inputs.shape[0] <= num_inducing:
-            self.allocator_used = self.__class__.__name__
-            return unique_inputs
-
-        # Otherwise, fall back to the provided allocator (e.g., KMeansAllocator)
-        if inputs.shape[0] <= num_inducing:
-            self.allocator_used = self.__class__.__name__
-            return inputs
-        else:
-            self.allocator_used = self.fallback_allocator.__class__.__name__
-            return self.fallback_allocator.allocate_inducing_points(
-                inputs=inputs,
-                covar_module=covar_module,
-                num_inducing=num_inducing,
-                input_batch_shape=input_batch_shape,
-            )
-
-    @classmethod
-    def get_config_options(
-        cls,
-        config: Config,
-        name: Optional[str] = None,
-        options: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, Any]:
-        """Get configuration options for the AutoAllocator.
-
-        Args:
-            config (Config): Configuration object.
-            name (str, optional): Name of the allocator, defaults to None.
-            options (Dict[str, Any], optional): Additional options, defaults to None.
-
-        Returns:
-            Dict[str, Any]: Configuration options for the AutoAllocator.
-        """
-        if name is None:
-            name = cls.__name__
-        lb = config.gettensor("common", "lb")
-        ub = config.gettensor("common", "ub")
-        bounds = torch.stack((lb, ub))
-        fallback_allocator_cls = config.getobj(
-            name, "fallback_allocator", fallback=KMeansAllocator
-        )
-        fallback_allocator = (
-            fallback_allocator_cls.from_config(config)
-            if hasattr(fallback_allocator_cls, "from_config")
-            else fallback_allocator_cls()
+        # Auto allocator actually just wraps the Kmeans allocator
+        allocator = KMeansAllocator(dim=self.dim)
+
+        points = allocator.allocate_inducing_points(
+            inputs=inputs,
+            covar_module=covar_module,
+            num_inducing=num_inducing,
+            input_batch_shape=input_batch_shape,
         )
 
-        return {"fallback_allocator": fallback_allocator, "bounds": bounds}
+        self.last_allocator_used = allocator.last_allocator_used
+        return points

aepsych/models/inducing_points/base.py

@@ -9,57 +9,30 @@
 class BaseAllocator(InducingPointAllocator, ConfigurableMixin):
     """Base class for inducing point allocators."""
 
-    def __init__(self, bounds: Optional[torch.Tensor] = None) -> None:
+    def __init__(self, dim: int, *args, **kwargs) -> None:
         """
         Initialize the allocator with optional bounds.
 
         Args:
-            bounds (torch.Tensor, optional): Bounds for allocating points. Should be of shape (2, d).
+            dim (int): Dimensionality of the search space.
+            *args, **kwargs: Other allocator specific arguments.
         """
-        self.bounds = bounds
-        self.dim = self._initialize_dim()
-
-    def _initialize_dim(self) -> Optional[int]:
-        """
-        Initialize the dimension `dim` based on the bounds, if available.
-
-        Returns:
-            int: The dimension `d` if bounds are provided, or None otherwise.
-        """
-        if self.bounds is not None:
-            # Validate bounds and extract dimension
-            assert self.bounds.shape[0] == 2, "Bounds must have shape (2, d)!"
-            lb, ub = self.bounds[0], self.bounds[1]
-            self.lb, self.ub = lb, ub
-            for i, (l, u) in enumerate(zip(lb, ub)):
-                assert (
-                    l <= u
-                ), f"Lower bound {l} is not less than or equal to upper bound {u} on dimension {i}!"
-            return self.bounds.shape[1]  # Number of dimensions (d)
-        return None
-
-    def _determine_dim_from_inputs(self, inputs: torch.Tensor) -> int:
-        """
-        Determine dimension `dim` from the inputs tensor.
-
-        Args:
-            inputs (torch.Tensor): Input tensor of shape (..., d).
-
-        Returns:
-            int: The inferred dimension `d`.
-        """
-        return inputs.shape[-1]
+        self.dim = dim
+        self.last_allocator_used: Optional[InducingPointAllocator] = None
 
     @abstractmethod
     def allocate_inducing_points(
         self,
-        inputs: Optional[torch.Tensor],
-        covar_module: Optional[torch.nn.Module],
-        num_inducing: int,
-        input_batch_shape: torch.Size,
+        inputs: Optional[torch.Tensor] = None,
+        covar_module: Optional[torch.nn.Module] = None,
+        num_inducing: int = 100,
+        input_batch_shape: torch.Size = torch.Size([]),
     ) -> torch.Tensor:
         """
-        Abstract method for allocating inducing points.
+        Abstract method for allocating inducing points. Must replace the
+        last_allocator_used attribute for what was actually used to produce the
+        inducing points. Dummy points should be made when it is not possible to create
+        inducing points (e.g., inputs is None).
 
         Args:
             inputs (torch.Tensor, optional): Input tensor, implementation-specific.
@@ -70,56 +43,18 @@ def allocate_inducing_points(
         Returns:
             torch.Tensor: Allocated inducing points.
         """
-        if self.dim is None and inputs is not None:
-            self.dim = self._determine_dim_from_inputs(inputs)
-
-        raise NotImplementedError("This method should be implemented by subclasses.")
-
-    @abstractmethod
-    def _get_quality_function(self) -> Optional[Any]:
-        """
-        Abstract method for returning a quality function if required.
-
-        Returns:
-            None or Callable: Quality function if needed.
-        """
-        raise NotImplementedError("This method should be implemented by subclasses.")
 
-
-class DummyAllocator(BaseAllocator):
-    def __init__(self, bounds: torch.Tensor) -> None:
-        """Initialize the DummyAllocator with bounds.
+    def _allocate_dummy_points(self, num_inducing: int = 100) -> torch.Tensor:
+        """Return dummy inducing points with the correct dimensionality.
 
         Args:
-            bounds (torch.Tensor): Bounds for allocating points. Should be of shape (2, d).
+            num_inducing (int): Number of inducing points to make, defaults to 100.
         """
-        super().__init__(bounds=bounds)
-        self.bounds: torch.Tensor = bounds
-
-    def _get_quality_function(self) -> None:
-        """DummyAllocator does not require a quality function, so this returns None."""
-        return None
-
-    def allocate_inducing_points(
-        self,
-        inputs: Optional[torch.Tensor] = None,
-        covar_module: Optional[torch.nn.Module] = None,
-        num_inducing: int = 10,
-        input_batch_shape: torch.Size = torch.Size([]),
-    ) -> torch.Tensor:
-        """Allocate inducing points by returning zeros of the appropriate shape.
+        self.last_allocator_used = None
+        return torch.zeros(num_inducing, self.dim)
 
-        Args:
-            inputs (torch.Tensor): Input tensor, not required for DummyAllocator.
-            covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here.
-            num_inducing (int, optional): The number of inducing points to generate. Defaults to 10.
-            input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here.
-
-        Returns:
-            torch.Tensor: A (num_inducing, d)-dimensional tensor of zeros.
-        """
-        self.allocator_used = self.__class__.__name__
-        return torch.zeros(num_inducing, self.bounds.shape[-1])
+    def _get_quality_function(self):
+        return super()._get_quality_function()
 
     @classmethod
     def get_config_options(
@@ -128,19 +63,20 @@ def get_config_options(
         name: Optional[str] = None,
         options: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
-        """Get configuration options for the DummyAllocator.
+        """Get configuration options for the allocator.
 
         Args:
             config (Config): Configuration object.
-            name (str, optional): Name of the allocator, defaults to None.
+            name (str, optional): Name of the allocator, defaults to None. Ignored.
             options (Dict[str, Any], optional): Additional options, defaults to None.
 
         Returns:
             Dict[str, Any]: Configuration options for the DummyAllocator.
         """
-        if name is None:
-            name = cls.__name__
-        lb = config.gettensor("common", "lb")
-        ub = config.gettensor("common", "ub")
-        bounds = torch.stack((lb, ub))
-        return {"bounds": bounds}
+        if options is None:
+            options = {}
+
+        if "dim" not in options:
+            options["dim"] = len(config.getlist("common", "parnames", element_type=str))
+
+        return options