diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index 07c22a27d..9fb05044b 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/sbi/inference/posteriors/vi_posterior.py @@ -20,7 +20,6 @@ if TYPE_CHECKING: from sbi.inference.posteriors.posterior_parameters import VIPosteriorParameters from sbi.inference.potentials.base_potential import BasePotential, CustomPotential -from sbi.neural_nets.estimators.base import ConditionalDensityEstimator from sbi.neural_nets.estimators.zuko_flow import ZukoUnconditionalFlow from sbi.neural_nets.factory import ZukoFlowType from sbi.neural_nets.net_builders.flow import ( @@ -32,6 +31,7 @@ from sbi.samplers.vi.vi_utils import ( AdaptedVariationalDistribution, LearnableGaussian, + TransformedConditionalZukoFlow, TransformedZukoFlow, check_variational_distribution, detach_and_deepcopy, @@ -46,6 +46,7 @@ from sbi.utils.torchutils import ( atleast_2d_float32_tensor, ensure_theta_batched, + move_all_tensor_to_device, process_device, ) from sbi.utils.user_input_checks_utils import move_distribution_to_device @@ -208,7 +209,7 @@ def __init__( self._mode: Optional[Literal["single_x", "amortized"]] = None # Amortized mode: conditional flow q(θ|x) - self._amortized_q: Optional[ConditionalDensityEstimator] = None + self._amortized_q: Optional[TransformedConditionalZukoFlow] = None self._num_transforms: int = num_transforms self._hidden_features: int = hidden_features @@ -278,9 +279,29 @@ def to(self, device: Union[str, torch.device]) -> "VIPosterior": self._q._link_transform = self.link_transform # type: ignore[union-attr] if self._amortized_q is not None: self._amortized_q.to(device) + # Keep the trained link: rebuilding it can re-estimate prior moments. + link_transform = self._amortized_q._link_transform + move_all_tensor_to_device(link_transform, device) + self._rebind_amortized_q(self._prior, link_transform) return self + @staticmethod + def _prior_support(prior: Distribution): + """Support of the prior, or None if it does not define one.""" + try: + return prior.support + except (AttributeError, NotImplementedError): + return None + + def _rebind_amortized_q( + self, prior: Distribution, link_transform: TorchTransform + ) -> None: + """Point the amortized flow's link and support at tensors on this device.""" + assert self._amortized_q is not None + self._amortized_q._link_transform = link_transform + self._amortized_q._support = self._prior_support(prior) + def _build_unconditional_flow( self, flow_type: str, @@ -411,7 +432,7 @@ def _build_conditional_flow( hidden_features: int = 32, z_score_theta: Literal["none", "independent", "structured"] = "independent", z_score_x: Literal["none", "independent", "structured"] = "independent", - ) -> ConditionalDensityEstimator: + ) -> TransformedConditionalZukoFlow: """Build a conditional Zuko flow for amortized variational inference. Args: @@ -427,7 +448,7 @@ def _build_conditional_flow( structured data like images. Returns: - ConditionalDensityEstimator: The constructed conditional flow q(θ|x). + TransformedConditionalZukoFlow: The conditional flow q(θ|x) in prior space. Raises: ValueError: If flow_type is not supported. @@ -442,14 +463,19 @@ def _build_conditional_flow( f"Supported types: {[t.name for t in ZukoFlowType]}." ) from e - return build_zuko_flow( + # The flow models θ in unconstrained space; the wrapper maps it to the prior + # support so that bounded priors are not evaluated outside their support. + flow = build_zuko_flow( flow_type.value.upper(), - batch_x=theta, # θ is what we model + batch_x=self.link_transform.inv(theta), # θ is what we model batch_y=x, # x is the condition z_score_x=z_score_theta, # z-score for θ (naming mismatch) z_score_y=z_score_x, # z-score for x condition num_transforms=num_transforms, hidden_features=hidden_features, + ) + return TransformedConditionalZukoFlow( + flow, self.link_transform, self._prior_support(self._prior) ).to(self._device) @property @@ -561,6 +587,17 @@ def set_q( self._amortized_q = deepcopy(getattr(q, "_amortized_q", None)) if self._amortized_q is not None: self._amortized_q.to(self._device) + # Link and support come from the source posterior's prior (not this + # one's), copied so that the source stays on its own device. + prior = move_distribution_to_device( + detach_and_deepcopy(q._prior), self._device + ) + if q.theta_transform is None: + theta_transform = mcmc_transform(prior, device=self._device) + else: + theta_transform = detach_and_deepcopy(q.theta_transform) + move_all_tensor_to_device(theta_transform, self._device) + self._rebind_amortized_q(prior, theta_transform.inv) self.vi_method = q.vi_method # type: ignore self._prior = q._prior self._x = q._x diff --git a/sbi/samplers/vi/vi_utils.py b/sbi/samplers/vi/vi_utils.py index d41a6a621..f27fc53d8 100644 --- a/sbi/samplers/vi/vi_utils.py +++ b/sbi/samplers/vi/vi_utils.py @@ -20,9 +20,10 @@ Normal, TransformedDistribution, ) +from torch.distributions.constraints import Constraint from torch.nn import Module -from sbi.neural_nets.estimators.zuko_flow import ZukoUnconditionalFlow +from sbi.neural_nets.estimators.zuko_flow import ZukoFlow, ZukoUnconditionalFlow from sbi.sbi_types import Shape, TorchTransform, VariationalDistribution from sbi.utils.torchutils import _base_recursor @@ -136,6 +137,71 @@ def sample_and_log_prob(self, sample_shape: Shape) -> tuple[Tensor, Tensor]: return theta, log_prob_theta +class TransformedConditionalZukoFlow(nn.Module): + """Conditional counterpart of `TransformedZukoFlow` for amortized VI. + + The wrapped flow q(z|x) lives in unconstrained space. Samples and log-probs of + this wrapper are in the prior's (constrained) space, so a bounded prior is never + evaluated outside its support. + """ + + def __init__( + self, + flow: ZukoFlow, + link_transform: TorchTransform, + support: Optional[Constraint] = None, + ): + """Initialize the wrapper. + + Args: + flow: Conditional Zuko flow operating in unconstrained space. + link_transform: Transform from unconstrained to constrained space. + support: Support of the prior. `log_prob` returns `-inf` outside of it + (the inverse link would otherwise clamp such points into it). + """ + super().__init__() + self._flow = flow + self._link_transform = link_transform + self._support = support + + @staticmethod + def _sum_log_det(log_det: Tensor, like: Tensor) -> Tensor: + # Elementwise transforms return per-dimension Jacobians; sum them if needed. + return log_det.sum(dim=-1) if log_det.dim() > like.dim() else log_det + + def sample(self, sample_shape: Shape, condition: Tensor) -> Tensor: + """Sample θ ~ q(θ|x) in constrained space.""" + theta = self._link_transform(self._flow.sample(sample_shape, condition)) + assert isinstance(theta, Tensor) + return theta + + def log_prob(self, theta: Tensor, condition: Tensor) -> Tensor: + """Log density of constrained θ, including the Jacobian of the link.""" + z = self._link_transform.inv(theta) + assert isinstance(z, Tensor) + log_prob_z = self._flow.log_prob(z, condition) + log_det = self._link_transform.inv.log_abs_det_jacobian(theta, z) + log_prob = log_prob_z + self._sum_log_det(log_det, log_prob_z) + if self._support is None: + return log_prob + in_support = self._support.check(theta) + if in_support.dim() > log_prob.dim(): + in_support = in_support.all(dim=-1) + return torch.where(in_support, log_prob, torch.full_like(log_prob, -torch.inf)) + + def sample_and_log_prob( + self, sample_shape: Shape, condition: Tensor + ) -> tuple[Tensor, Tensor]: + """Reparameterized constrained samples and their log density.""" + z, log_prob_z = self._flow.sample_and_log_prob( + torch.Size(sample_shape), condition + ) + theta = self._link_transform(z) + assert isinstance(theta, Tensor) + log_det = self._link_transform.log_abs_det_jacobian(z, theta) + return theta, log_prob_z - self._sum_log_det(log_det, log_prob_z) + + class LearnableGaussian(nn.Module): """Learnable Gaussian distribution for variational inference. diff --git a/tests/vi_test.py b/tests/vi_test.py index 18b23cc03..02c819366 100644 --- a/tests/vi_test.py +++ b/tests/vi_test.py @@ -37,7 +37,7 @@ linear_gaussian, true_posterior_linear_gaussian_mvn_prior, ) -from sbi.utils import MultipleIndependent +from sbi.utils import BoxUniform, MultipleIndependent from sbi.utils.metrics import c2st, check_c2st from sbi.utils.user_input_checks import process_x @@ -878,6 +878,25 @@ def test_amortized_vi_batched_sampling(linear_gaussian_setup): assert samples.shape == (num_samples, num_obs, setup["num_dim"]) +def test_amortized_vi_with_bounded_prior(): + """Amortized VI trains and samples inside the support of a `BoxUniform` prior.""" + prior = BoxUniform(-3 * ones(2), 3 * ones(2)) + theta = prior.sample((300,)) + x = theta + 0.3 * torch.randn_like(theta) + nle = NLE(prior, show_progress_bars=False) + nle.append_simulations(theta, x).train(max_num_epochs=3) + posterior = nle.build_posterior(sample_with="vi") + + posterior.train_amortized(theta, x, max_num_iters=5, show_progress_bar=False) + + samples = posterior.sample((50,), x=x[:1], show_progress_bars=False) + assert torch.isfinite(samples).all() + assert (samples.abs() <= 3).all() + assert torch.isfinite(posterior.log_prob(samples, x=x[:1])).all() + outside = torch.tensor([[4.0, 0.0]]) + assert posterior.log_prob(outside, x=x[:1]).item() == -torch.inf + + @pytest.mark.slow def test_amortized_vi_log_prob(linear_gaussian_setup): """Test log_prob evaluation in amortized mode."""