[GENERAL SUPPORT]: How do I understand the `pruning_target_parameterization`

[zhqrbitee]

Question

Hi, I have an optimization case like below. It has multi objectives, which depends on a common set of parameters, but each also depends on individual parameters.

I'm seeking some way to tell the optimizer which objective depends on which parameters (actually we know it in advance), or config the optimizer so it can automatically prune unrelated parameters for an objective, so that it could improve the sampling efficiency.

I notice that the AX has introduced a 'simplify_parameter_changes' argument in latest version, which seems may work for my case. However, to enable it, it asks me to set the pruning_target_parameterization on optimization_config. There are some comments in the codes like below:

            pruning_target_parameterization: Arm containing the target values for
                irrelevant parameters. The target values are used to prune irrelevant
                parameters from candidates generated via Bayesian optimization: when
                Ax considers arms to suggest for the next trial, it will seek to have
                the proposed arms differ from the target arm as little as possible,
                without a loss in optimization performance. I.e. when suggested arms
                include parameter values that differ from the corresponding value in
                the target arm, the pruning methodology will check if that difference
                is expected to be meaningful w.r.t. the performance of the arm in
                consideration, and if not, the parameter value will be replaced with
                the corresponding value in the target arm.

But to be honest, I didn't get the meaning of target values for irrelevant parameters. So is there any more detailed explanation or example for it?

Please provide any relevant code snippet if applicable.

Code of Conduct

• I agree to follow this Ax's Code of Conduct

[mgrange1998]

This method in acquisition.py shows where the pruning takes place

Ax/ax/generators/torch/botorch_modular/acquisition.py

Lines 782 to 924 in 89b35c7

	def _prune_irrelevant_parameters(
	self,
	candidates: Tensor,
	search_space_digest: SearchSpaceDigest,
	inequality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
	fixed_features: dict[int, float] | None = None,
	) -> tuple[Tensor, Tensor]:
	r"""Prune irrelevant parameters from the candidates.
	The method involves first optimizing the AF without any notion of irrelevance.
	Then, the irrelevant parameters are pruned via a sequential greedy algorithm
	that sets each dimension to the target point value, and recomputes the AF value.
	The dimension that corresponds to the smallest reduction in AF value compared to
	the AF value of the dense original candidate is then set to the target value,
	and the algorithm proceeds until removing any dimension causes the reduction in
	the AF value of the proposed point to be less than some tolerance, relative to
	the original dense point.
	For `q > 1` cases, this is repeated for each point in the batch, after adding
	the previously pruned points as pending points. The incremental AF value is
	used in the stopping rule.
	Args:
	candidates: A `q x d`-dim Tensor of candidates.
	search_space_digest: The SearchSpaceDigest.
	inequality_constraints: A list of tuples (indices, coefficients, rhs),
	with each tuple encoding an inequality constraint of the form
	`\sum_i (X[indices[i]] * coefficients[i]) >= rhs`. `indices` and
	`coefficients` should be torch tensors. See the docstring of
	`make_scipy_linear_constraints` for an example. When q=1, or when
	applying the same constraint to each candidate in the batch
	(intra-point constraint), `indices` should be a 1-d tensor.
	For inter-point constraints, in which the constraint is applied to the
	whole batch of candidates, `indices` must be a 2-d tensor, where
	in each row `indices[i] =(k_i, l_i)` the first index `k_i` corresponds
	to the `k_i`-th element of the `q`-batch and the second index `l_i`
	corresponds to the `l_i`-th feature of that element.
	fixed_features: A map `{feature_index: value}` for features that
	should be fixed to a particular value during generation.
	Returns:
	A two-element tuple containing an `q x d`-dim tensor of generated
	candidates and a `q`-dim tensor with the associated acquisition values.
	"""
	target_point = none_throws(
	self._pruning_target_point,
	message="Must specify pruning_target_point to prune irrelevant parameters",
	)
	irrelevance_pruning_rtol = self.options.get("irrelevance_pruning_rtol", 0.1)
	initial_X_pending = self.X_pending
	pruned_af_vals = []
	excluded_indices = set(search_space_digest.fidelity_features)
	excluded_indices.update(set(search_space_digest.task_features))
	if fixed_features is not None:
	excluded_indices.update(set(fixed_features.keys()))
	orig_acqf = self.acqf
	# iterate over points in the batch and prune each point
	initial_active_dims = (candidates != target_point).sum(dim=-1)
	for i in range(candidates.shape[0]):
	# condition on previous pruned points
	self._condition_on_prev_candidates(
	prev_candidates=candidates[:i], initial_X_pending=initial_X_pending
	)
	# compute the AF value for the previous candidates so that we can compute
	# the incremental EI
	base_af_val = self._compute_baseline_acqf_value(
	last_candidate=candidates[i - 1 : i], fixed_features=fixed_features
	)
	# compute the AF value of the dense candidate (before pruning)
	with torch.no_grad():
	dense_af_val = self.evaluate(X=candidates[i : i + 1])
	# compute the incremental AF value of the dense candidate, relative to the
	# base AF value
	dense_incremental_af_val = (
	((dense_af_val.exp() if self.acqf._log else dense_af_val) - base_af_val)
	.clamp_min(0.0)
	.item()
	)
	# In the base case that no parameters are pruned, the final AF val is the
	# dense AF val
	final_af_val = dense_af_val
	# If the current incremental AF value is zero, then we skip pruning
	if dense_incremental_af_val > 0.0:
	remaining_indices = set(range(candidates.shape[-1])) - excluded_indices
	# remove features that are already set to target_point
	remaining_indices -= set(
	(candidates[i] == target_point).nonzero().view(-1).tolist()
	)
	# len(remaining_indices) - 1 is used here so that we do not prune
	# every dimension
	for _ in range(len(remaining_indices) - 1):
	indices = torch.tensor(
	list(remaining_indices),
	dtype=torch.long,
	device=candidates.device,
	)
	# create a `b x 1 x d`-dim tensor, where each batch has the
	# original candidate with a different dimension set to the
	# corresponding target value
	pruned_candidates = _expand_and_set_single_feature_to_target(
	X=candidates[i : i + 1],
	indices=indices,
	targets=target_point[indices],
	)
	# remove candidates that violate constraints after pruning
	pruned_candidates, indices = _remove_infeasible_candidates(
	candidates=pruned_candidates,
	indices=indices,
	inequality_constraints=inequality_constraints,
	)
	if pruned_candidates.shape[0] == 0:
	# no feasible points, continue to
	# next candidate
	break
	# determine which dimension to prune based on the reduction
	# in incremental AF value relative to original dense candidate
	min_idx, best_af_val = self._get_best_pruned_candidate(
	pruned_candidates=pruned_candidates,
	irrelevance_pruning_rtol=irrelevance_pruning_rtol,
	dense_incremental_af_val=dense_incremental_af_val,
	base_af_val=base_af_val,
	)
	if min_idx is not None:
	candidates[i] = pruned_candidates[min_idx]
	remaining_indices.remove(
	int(none_throws(indices)[min_idx].item())
	)
	final_af_val = best_af_val
	else:
	# if min_idx is None that means that no candidate met
	# the relative tolerance threshold, so we stop pruning
	# this candidate
	break
	pruned_af_vals.append(final_af_val)
	self.acqf = orig_acqf
	# store the number of pruned dimensions for each candidate to return
	# this via gen_metadata
	self.num_pruned_dims = (
	initial_active_dims - (candidates != target_point).sum(dim=-1)
	).tolist()
	return candidates, torch.cat(pruned_af_vals, dim=0)

"""
Prune irrelevant parameters from the candidates.

The method involves first optimizing the AF without any notion of irrelevance.
Then, the irrelevant parameters are pruned via a sequential greedy algorithm
that sets each dimension to the target point value, and recomputes the AF value.
The dimension that corresponds to the smallest reduction in AF value compared to
the AF value of the dense original candidate is then set to the target value,
and the algorithm proceeds until removing any dimension causes the reduction in
the AF value of the proposed point to be less than some tolerance, relative to
the original dense point.

For q > 1 cases, this is repeated for each point in the batch, after adding
the previously pruned points as pending points. The incremental AF value is
used in the stopping rule.
"""

At a high level:
pruning_target_parameterization specifies "target values" for "irrelevant parameters".
The pruning process creates pruned candidates by replacing the value of "irrelevant parameter"s with the specified "target value" with a sequential greedy algorithm.
Then, in "get_best_pruned_candidate" the candidate with the best performance for the acquisition function (AF) value is selected.

cc @sdaulton for more context on pruning_target_parameterization

cc @hvarfner for guidance on specifying relationships between individual parameters and associated objectives

One thing you could try is setting your N parameters as irrelevant parameters with a target value you know in advance. Then during the acquisition phase, each of the N parameters will consider pruned variants with the parameter value set to its target if the decrease to acquisition function value is low.

[hvarfner]

Hi @zhqrbitee ,

Interesting problem! I would say that we don't have support for that exact setup, but running a SaasFullyBayesianSingleTaskGP (a surrogate that is specifically designed to identify irrelevant dimensions, per objective) and trusting the model to figure out the irrelevant dimensions would be a low-effort solution that would most likely get the job done.

Getting exactly what you are after would entail specifying unique models per objective using a ModelListGP botorch_model_class and altering the active_dims at the kernel level of the surrogate. This would be very tedious, and almost certainly not be worth the effort.