sabr

Select adversarial input subregions and compute SABR training bounds.

bound_sabr(hardened_model, original_model, data, target, eps, subselection_ratio, device='cuda', n_classes=10, x_L=None, x_U=None, data_min=None, data_max=None, n_steps=8, step_size=0.5, restarts=1, early_stopping=True, intermediate_bound_model=None, decay_factor=0.1, decay_checkpoints=(4, 7), return_adv_output=False, pgd_ptb=None, relu_upper_retention=1.0, pgd_eps_factor=1.0)

Compute training bounds over an adversarially selected input subregion.

SABR uses PGD to select a difficult point, then propagates a smaller, contained region around that point using IBP. These bounds describe the selected subregion and do not certify the full nominal perturbation region.

Parameters:

Name	Type	Description	Default
hardened_model	BoundedModule	The auto_LiRPA model.	required
original_model	Module	The original neural network model.	required
data	Tensor	The input data tensor.	required
target	Tensor	The target labels tensor.	required
eps	float	Perturbation radius in model input coordinates.	required
subselection_ratio	float	The ratio for subselection of the epsilon for the IBP bounding during SABR.	required
device	str	The device to run the computation on.	'cuda'
n_classes	int	The number of classes for classification.	10
x_L	Tensor	The lower bound of the input data.	None
x_U	Tensor	The upper bound of the input data.	None
data_min	Tensor	The minimum value of the input data.	None
data_max	Tensor	The maximum value of the input data.	None
n_steps	int	The number of steps for the attack.	8
step_size	float	The step size for the attack.	0.5
restarts	int	The number of restarts for the attack.	1
early_stopping	bool	Whether to use early stopping.	True
intermediate_bound_model	Module	The intermediate bound model. If provided, the SABR bounds of the intermediate model are returned. This is needed during STAPS bound calculation.	None
decay_factor	float	The decay factor for the attack.	0.1
decay_checkpoints	tuple	The decay checkpoints for the attack.	(4, 7)
return_adv_output	bool	Whether to return the adversarial output.	False
pgd_ptb	PerturbationLpNorm	Alternate bounds for the PGD centre search. The resulting SABR box is still clamped to the nominal region described by eps or x_L/x_U.	None
relu_upper_retention	float	Unstable ReLU upper bound fraction for this training bound; default 1. Values below 1 are not sound certificates.	1.0
pgd_eps_factor	float	Positive PGD search-range factor; nominal IBP bounds are unchanged. Explicit pgd_ptb, when provided, takes precedence.	1.0

Returns:

Type	Description
Tuple[Tensor, Tensor, Tensor]	The lower and upper bounds of the model's output, and the adversarial output if return_adv_output is True.

Source code in CTRAIN/bound/sabr.py

def bound_sabr(hardened_model, original_model, data, target, eps, subselection_ratio, device='cuda', n_classes=10, x_L=None, x_U=None, data_min=None, data_max=None, n_steps=8, step_size=.5, restarts=1, early_stopping=True, intermediate_bound_model=None, decay_factor=0.1, decay_checkpoints=(4,7), return_adv_output=False, pgd_ptb=None, relu_upper_retention=1.0, pgd_eps_factor=1.0):
    """Compute training bounds over an adversarially selected input subregion.

    SABR uses PGD to select a difficult point, then propagates a smaller,
    contained region around that point using IBP. These bounds describe the
    selected subregion and do not certify the full nominal perturbation region.

    Args:
        hardened_model (auto_LiRPA.BoundedModule): The auto_LiRPA model.
        original_model (torch.nn.Module): The original neural network model.
        data (Tensor): The input data tensor.
        target (Tensor): The target labels tensor.
        eps (float): Perturbation radius in model input coordinates.
        subselection_ratio (float): The ratio for subselection of the epsilon for the IBP
            bounding during SABR.
        device (str, optional): The device to run the computation on.
        n_classes (int, optional): The number of classes for classification.
        x_L (Tensor, optional): The lower bound of the input data.
        x_U (Tensor, optional): The upper bound of the input data.
        data_min (Tensor, optional): The minimum value of the input data.
        data_max (Tensor, optional): The maximum value of the input data.
        n_steps (int, optional): The number of steps for the attack.
        step_size (float, optional): The step size for the attack.
        restarts (int, optional): The number of restarts for the attack.
        early_stopping (bool, optional): Whether to use early stopping.
        intermediate_bound_model (torch.nn.Module, optional): The intermediate bound model.
            If provided, the SABR bounds
            of the intermediate model are
            returned. This is needed
            during STAPS bound
            calculation.
        decay_factor (float, optional): The decay factor for the attack.
        decay_checkpoints (tuple, optional): The decay checkpoints for the attack.
        return_adv_output (bool, optional): Whether to return the adversarial output.
        pgd_ptb (PerturbationLpNorm, optional): Alternate bounds for the PGD
            centre search. The resulting SABR box is still clamped to the
            nominal region described by ``eps`` or ``x_L``/``x_U``.
        relu_upper_retention (float): Unstable ReLU upper bound fraction for this training
            bound; default 1. Values below 1 are not sound
            certificates.
        pgd_eps_factor (float): Positive PGD search-range factor; nominal IBP bounds are
            unchanged. Explicit pgd_ptb, when provided, takes
            precedence.

    Returns:
        (Tuple[Tensor, Tensor, Tensor]): The lower and upper bounds of the model's output, and the adversarial output if return_adv_output is True.
    """
    validate_hybrid_options(relu_upper_retention, pgd_eps_factor)
    if pgd_ptb is None and pgd_eps_factor != 1:
        nominal = PerturbationLpNorm(
            norm=np.inf, eps=eps if eps is not None else np.inf,
            x_L=x_L if x_L is not None else torch.clamp(data - eps, data_min, data_max),
            x_U=x_U if x_U is not None else torch.clamp(data + eps, data_min, data_max))
        pgd_ptb = get_pgd_ptb(nominal, data, pgd_eps_factor, data_min, data_max, eps=eps)
    models = (hardened_model, original_model) + (() if intermediate_bound_model is None else (intermediate_bound_model,))
    with preserve_model_state(*models):
        with preserve_model_state(hardened_model, original_model):
            hardened_model.eval()
            original_model.eval()

            propagation_inputs, tau, x_adv = get_propagation_region(
                model=hardened_model,
                data=data,
                data_min=data_min,
                data_max=data_max,
                target=target,
                eps=eps if (x_L is None and x_U is None) else None,
                subselection_ratio=subselection_ratio,
                n_steps=n_steps,
                step_size=step_size,
                restarts=restarts,
                early_stopping=early_stopping,
                x_L=x_L,
                x_U=x_U,
                pgd_x_L=None if pgd_ptb is None else pgd_ptb.x_L,
                pgd_x_U=None if pgd_ptb is None else pgd_ptb.x_U,
                decay_checkpoints=decay_checkpoints,
                decay_factor=decay_factor
            )


        ptb = PerturbationLpNorm(
            eps=tau,
            norm=np.inf,
            x_L=torch.clamp(propagation_inputs - tau, data_min, data_max).to(device),
            x_U=torch.clamp(propagation_inputs + tau, data_min, data_max).to(device)
        )

        # Pass input through network to set batch statistics
        adv_output = hardened_model(x_adv)

        # Use intermediate_bound_model if provided and return intermediate bounds (as needed by STAPS), otherwise use hardened_model
        lb, ub = bound_ibp(
            model=hardened_model if intermediate_bound_model is None else intermediate_bound_model,
            ptb=ptb,
            data=propagation_inputs,
            # Only provide target if intermediate_bound_model is not used (as we are not interested in final bound margins)
            target=target if intermediate_bound_model is None else None,
            n_classes=n_classes,
            bound_upper=True,
            reuse_input=False,
            relu_upper_retention=relu_upper_retention
        )

        if return_adv_output:
            return lb, ub, adv_output
        return lb, ub

get_propagation_region(model, data, target, subselection_ratio, step_size, n_steps, restarts, x_L=None, x_U=None, data_min=None, data_max=None, eps=None, early_stopping=True, decay_factor=0.1, decay_checkpoints=(4, 7), pgd_x_L=None, pgd_x_U=None, *, loss_fn=None, is_success=None)

Select a contained SABR subregion around a PGD-generated point.

The half-width in each coordinate is subselection_ratio times the nominal interval half-width. Clipping the centre keeps this subregion inside the nominal box, including at input boundaries. Separate PGD endpoints may widen the search without enlarging the region used for propagation.

Parameters:

Name	Type	Description	Default
model	Module	The neural network model.	required
data	Tensor	The input data tensor.	required
target	Tensor	The target labels tensor.	required
subselection_ratio	float	The ratio for subselection.	required
step_size	float	The step size for the attack.	required
n_steps	int	The number of steps for the attack.	required
restarts	int	The number of restarts for the attack.	required
x_L	Tensor	The lower bound of the input data.	None
x_U	Tensor	The upper bound of the input data.	None
data_min	Tensor	The minimum value of the input data.	None
data_max	Tensor	The maximum value of the input data.	None
eps	float	Perturbation radius in model input coordinates.	None
early_stopping	bool	Whether to use early stopping.	True
decay_factor	float	The decay factor for the attack.	0.1
decay_checkpoints	tuple	The decay checkpoints for the attack.	(4, 7)
pgd_x_L	Tensor	Alternate lower bound for the PGD search.	None
pgd_x_U	Tensor	Alternate upper bound for the PGD search.	None
loss_fn	callable	Optional per-example attack objective to maximise, called as loss_fn(output, target).	None
is_success	callable	Optional per-example Boolean counterexample test, called as is_success(output, target).	None

Returns:

Type	Description
Tuple[Tensor, float, Tensor]	The propagation inputs, tau value, and adversarial examples.

Source code in CTRAIN/bound/sabr.py

def get_propagation_region(model, data, target, subselection_ratio, step_size, n_steps, restarts, x_L=None, x_U=None, data_min=None, data_max=None, eps=None, early_stopping=True, decay_factor=.1, decay_checkpoints=(4, 7), pgd_x_L=None, pgd_x_U=None, *, loss_fn=None, is_success=None):
    """Select a contained SABR subregion around a PGD-generated point.

    The half-width in each coordinate is subselection_ratio times the nominal
    interval half-width. Clipping the centre keeps this subregion inside the
    nominal box, including at input boundaries. Separate PGD endpoints may widen
    the search without enlarging the region used for propagation.

    Args:
        model (torch.nn.Module): The neural network model.
        data (Tensor): The input data tensor.
        target (Tensor): The target labels tensor.
        subselection_ratio (float): The ratio for subselection.
        step_size (float): The step size for the attack.
        n_steps (int): The number of steps for the attack.
        restarts (int): The number of restarts for the attack.
        x_L (Tensor, optional): The lower bound of the input data.
        x_U (Tensor, optional): The upper bound of the input data.
        data_min (Tensor, optional): The minimum value of the input data.
        data_max (Tensor, optional): The maximum value of the input data.
        eps (float, optional): Perturbation radius in model input coordinates.
        early_stopping (bool, optional): Whether to use early stopping.
        decay_factor (float, optional): The decay factor for the attack.
        decay_checkpoints (tuple, optional): The decay checkpoints for the attack.
        pgd_x_L (Tensor, optional): Alternate lower bound for the PGD search.
        pgd_x_U (Tensor, optional): Alternate upper bound for the PGD search.
        loss_fn (callable, optional): Optional per-example attack objective to maximise, called as
            loss_fn(output, target).
        is_success (callable, optional): Optional per-example Boolean counterexample test, called as
            is_success(output, target).

    Returns:
        (Tuple[Tensor, float, Tensor]): The propagation inputs, tau value, and adversarial examples.
    """
    device = data.device if data is not None else x_L.device
    if not (
        (x_L is None and x_U is None and eps is not None)
        or (x_L is not None and x_U is not None and eps is None)
    ):
        raise ValueError("Provide either eps or x_L/x_U, but not both")
    if (pgd_x_L is None) != (pgd_x_U is None):
        raise ValueError("Provide both PGD bounds or neither")
    if eps is not None and data is not None:
        x_L=torch.clamp(data - eps, data_min, data_max).to(device)
        x_U=torch.clamp(data + eps, data_min, data_max).to(device)
    if not 0 < subselection_ratio <= 1:
        raise ValueError("subselection_ratio must be in (0, 1]")
    if torch.any(x_L > x_U):
        raise ValueError("Nominal lower bounds must not exceed upper bounds")
    # Shrink each clipped interval, including pixels on the image boundary.
    tau = subselection_ratio * (x_U - x_L) / 2

    attack_x_L = x_L if pgd_x_L is None else pgd_x_L
    attack_x_U = x_U if pgd_x_U is None else pgd_x_U
    if attack_x_L.shape != x_L.shape or attack_x_U.shape != x_U.shape:
        raise ValueError("PGD bounds must have the same shape as the nominal bounds")
    if torch.any(attack_x_L > attack_x_U):
        raise ValueError("PGD lower bounds must not exceed upper bounds")

    attack_options = {} if loss_fn is None and is_success is None else dict(loss_fn=loss_fn, is_success=is_success)
    x_adv = pgd_attack(
        model=model,
        data=data,
        target=target,
        x_L=attack_x_L,
        x_U=attack_x_U,
        n_steps=n_steps,
        step_size=step_size,
        restarts=restarts,
        early_stopping=early_stopping,
        device=device,
        decay_checkpoints=decay_checkpoints,
        decay_factor=decay_factor,
        **attack_options
    )

    propagation_inputs = torch.clamp(x_adv, x_L + tau, x_U - tau) # called midpoints in SABR code
    tau = torch.as_tensor(tau, device=device)
    return propagation_inputs, tau, x_adv