taps

Combine feature bounds, latent searches and gradient links for TAPS and STAPS training.

GradExpander

Bases: Function

Scale a TAPS loss gradient while preserving its forward value.

This custom autograd operation lets the training objective weight the latent-loss derivative by alpha without multiplying the reported loss.

Source code in CTRAIN/bound/taps.py

class GradExpander(torch.autograd.Function):
    """Scale a TAPS loss gradient while preserving its forward value.

    This custom autograd operation lets the training objective weight the
    latent-loss derivative by alpha without multiplying the reported loss.
    """

    @staticmethod
    def forward(ctx, x, alpha:float=1):
        """Return x unchanged and retain the requested gradient multiplier.

        Args:
            ctx (torch.autograd.function.FunctionCtx): The context object that can be used to stash information
                for backward computation.
            x (torch.Tensor): The input tensor.
            alpha (float, optional): A scaling factor.

        Returns:
            (torch.Tensor): The input tensor `x`.
        """
        ctx.alpha = alpha
        return x

    @staticmethod
    def backward(ctx, grad_x):
        """Multiply the incoming gradient by alpha; return no gradient for the multiplier.

        Args:
            ctx (torch.autograd.function.FunctionCtx): The context object that can be used to stash information for backward
                computation.
            grad_x (torch.Tensor): The gradient of the loss with respect to the output of the forward pass.

        Returns:
            (Tuple[Tensor, None]): A tuple containing the gradient of the loss with respect to the input of the forward pass and None (as there is no gradient with respect to the second input).
        """
        return ctx.alpha * grad_x, None

backward(ctx, grad_x) staticmethod

Multiply the incoming gradient by alpha; return no gradient for the multiplier.

Parameters:

Name	Type	Description	Default
ctx	FunctionCtx	The context object that can be used to stash information for backward computation.	required
grad_x	Tensor	The gradient of the loss with respect to the output of the forward pass.	required

Returns:

Type	Description
Tuple[Tensor, None]	A tuple containing the gradient of the loss with respect to the input of the forward pass and None (as there is no gradient with respect to the second input).

Source code in CTRAIN/bound/taps.py

@staticmethod
def backward(ctx, grad_x):
    """Multiply the incoming gradient by alpha; return no gradient for the multiplier.

    Args:
        ctx (torch.autograd.function.FunctionCtx): The context object that can be used to stash information for backward
            computation.
        grad_x (torch.Tensor): The gradient of the loss with respect to the output of the forward pass.

    Returns:
        (Tuple[Tensor, None]): A tuple containing the gradient of the loss with respect to the input of the forward pass and None (as there is no gradient with respect to the second input).
    """
    return ctx.alpha * grad_x, None

forward(ctx, x, alpha=1) staticmethod

Return x unchanged and retain the requested gradient multiplier.

Parameters:

Name	Type	Description	Default
ctx	FunctionCtx	The context object that can be used to stash information for backward computation.	required
x	Tensor	The input tensor.	required
alpha	float	A scaling factor.	1

Returns:

Type	Description
Tensor	The input tensor x.

Source code in CTRAIN/bound/taps.py

@staticmethod
def forward(ctx, x, alpha:float=1):
    """Return x unchanged and retain the requested gradient multiplier.

    Args:
        ctx (torch.autograd.function.FunctionCtx): The context object that can be used to stash information
            for backward computation.
        x (torch.Tensor): The input tensor.
        alpha (float, optional): A scaling factor.

    Returns:
        (torch.Tensor): The input tensor `x`.
    """
    ctx.alpha = alpha
    return x

RectifiedLinearGradientLink

Bases: Function

Link latent pivotal-point gradients to feature interval endpoints for TAPS.

The forward computation returns the selected point. Backward assigns endpoint gradients according to its distance from each interval boundary. This is the method's surrogate gradient, not the derivative of the identity forward computation. The threshold c determines the linked boundary region, and tol limits division near zero-width intervals.

Source code in CTRAIN/bound/taps.py

class RectifiedLinearGradientLink(torch.autograd.Function):
    """Link latent pivotal-point gradients to feature interval endpoints for TAPS.

    The forward computation returns the selected point. Backward assigns
    endpoint gradients according to its distance from each interval boundary.
    This is the method's surrogate gradient, not the derivative of the identity
    forward computation. The threshold c determines the linked boundary region,
    and tol limits division near zero-width intervals.
    """
    @staticmethod
    def forward(ctx, lb, ub, x, c:float, tol:float):
        """Return the pivotal point and retain the feature bounds for the surrogate backward pass.

        Args:
            ctx (torch.autograd.function.FunctionCtx): Context object to save information for backward computation.
            lb (torch.Tensor): Lower bound tensor.
            ub (torch.Tensor): Upper bound tensor.
            x (torch.Tensor): Input tensor.
            c (float): A constant used to determine the slope.
            tol (float): A tolerance value to avoid division by zero.

        Returns:
            (Tensor): The input tensor x.
        """
        ctx.save_for_backward(lb, ub, x)
        ctx.c = c
        ctx.tol = tol
        return x
    @staticmethod
    def backward(ctx, grad_x):
        """Distribute pivotal-point gradients to nearby lower and upper feature bounds.

        Args:
            ctx (torch.autograd.function.FunctionCtx): Context object containing saved tensors and constants.
            grad_x (torch.Tensor): Gradient of the loss with respect to the output of the forward function.

        Returns:
            (Tuple[Tensor, Tensor, None, None, None]): Gradients with respect to lb, ub, and None for other inputs.
        """
        lb, ub, x = ctx.saved_tensors
        c, tol = ctx.c, ctx.tol
        slackness = c * (ub - lb)
        # handle grad w.r.t. ub
        thre = (ub - slackness)
        Rectifiedgrad_mask = (x >= thre)
        grad_ub = (Rectifiedgrad_mask * grad_x * (x - thre).clamp(min=0.5*tol) / slackness.clamp(min=tol)).sum(dim=0, keepdim=True)
        # handle grad w.r.t. lb
        thre = (lb + slackness)
        Rectifiedgrad_mask = (x <= thre)
        grad_lb = (Rectifiedgrad_mask * grad_x * (thre - x).clamp(min=0.5*tol) / slackness.clamp(min=tol)).sum(dim=0, keepdim=True)
        # we don't need grad w.r.t. x and param
        return grad_lb, grad_ub, None, None, None

backward(ctx, grad_x) staticmethod

Distribute pivotal-point gradients to nearby lower and upper feature bounds.

Parameters:

Name	Type	Description	Default
ctx	FunctionCtx	Context object containing saved tensors and constants.	required
grad_x	Tensor	Gradient of the loss with respect to the output of the forward function.	required

Returns:

Type	Description
Tuple[Tensor, Tensor, None, None, None]	Gradients with respect to lb, ub, and None for other inputs.

Source code in CTRAIN/bound/taps.py

@staticmethod
def backward(ctx, grad_x):
    """Distribute pivotal-point gradients to nearby lower and upper feature bounds.

    Args:
        ctx (torch.autograd.function.FunctionCtx): Context object containing saved tensors and constants.
        grad_x (torch.Tensor): Gradient of the loss with respect to the output of the forward function.

    Returns:
        (Tuple[Tensor, Tensor, None, None, None]): Gradients with respect to lb, ub, and None for other inputs.
    """
    lb, ub, x = ctx.saved_tensors
    c, tol = ctx.c, ctx.tol
    slackness = c * (ub - lb)
    # handle grad w.r.t. ub
    thre = (ub - slackness)
    Rectifiedgrad_mask = (x >= thre)
    grad_ub = (Rectifiedgrad_mask * grad_x * (x - thre).clamp(min=0.5*tol) / slackness.clamp(min=tol)).sum(dim=0, keepdim=True)
    # handle grad w.r.t. lb
    thre = (lb + slackness)
    Rectifiedgrad_mask = (x <= thre)
    grad_lb = (Rectifiedgrad_mask * grad_x * (thre - x).clamp(min=0.5*tol) / slackness.clamp(min=tol)).sum(dim=0, keepdim=True)
    # we don't need grad w.r.t. x and param
    return grad_lb, grad_ub, None, None, None

forward(ctx, lb, ub, x, c, tol) staticmethod

Return the pivotal point and retain the feature bounds for the surrogate backward pass.

Parameters:

Name	Type	Description	Default
ctx	FunctionCtx	Context object to save information for backward computation.	required
lb	Tensor	Lower bound tensor.	required
ub	Tensor	Upper bound tensor.	required
x	Tensor	Input tensor.	required
c	float	A constant used to determine the slope.	required
tol	float	A tolerance value to avoid division by zero.	required

Returns:

Type	Description
Tensor	The input tensor x.

Source code in CTRAIN/bound/taps.py

@staticmethod
def forward(ctx, lb, ub, x, c:float, tol:float):
    """Return the pivotal point and retain the feature bounds for the surrogate backward pass.

    Args:
        ctx (torch.autograd.function.FunctionCtx): Context object to save information for backward computation.
        lb (torch.Tensor): Lower bound tensor.
        ub (torch.Tensor): Upper bound tensor.
        x (torch.Tensor): Input tensor.
        c (float): A constant used to determine the slope.
        tol (float): A tolerance value to avoid division by zero.

    Returns:
        (Tensor): The input tensor x.
    """
    ctx.save_for_backward(lb, ub, x)
    ctx.c = c
    ctx.tol = tol
    return x

bound_taps(original_model, hardened_model, bounded_blocks, data, target, n_classes, ptb, device='cuda', pgd_steps=20, pgd_restarts=1, pgd_step_size=0.2, pgd_decay_factor=0.2, pgd_decay_checkpoints=(5, 7), gradient_link_thresh=0.5, gradient_link_tolerance=1e-05, propagation='IBP', sabr_args=None, relu_upper_retention=1.0, pgd_eps_factor=1.0)

Estimate classifier margins using interval propagation and latent PGD.

TAPS propagates the input region through the feature extractor using IBP, then searches for pivotal points in its output interval. STAPS uses SABR for feature propagation. The returned latent margins are training proxies; use ordinary full-region bounds for sound certification.

Parameters:

Name	Type	Description	Default
original_model	Module	The original neural network model.	required
hardened_model	BoundedModule	The auto_LiRPA model.	required
bounded_blocks	list	List of bounded blocks of the model.	required
data	Tensor	The input data tensor.	required
target	Tensor	The target labels tensor.	required
n_classes	int	The number of classes for classification.	required
ptb	PerturbationLpNorm	The perturbation object defining the perturbation set.	required
device	str	The device to run the computation on.	'cuda'
pgd_steps	int	The number of steps for the PGD attack.	20
pgd_restarts	int	The number of restarts for the PGD attack.	1
pgd_step_size	float	Latent-PGD step-size factor; None uses max(0.2, 2 / steps) for positive step counts.	0.2
pgd_decay_factor	float	The decay factor for the PGD attack.	0.2
pgd_decay_checkpoints	tuple	The decay checkpoints for the PGD attack.	(5, 7)
gradient_link_thresh	float	The threshold for gradient linking.	0.5
gradient_link_tolerance	float	The tolerance for gradient linking.	1e-05
propagation	str	The propagation method to use ('IBP' or 'SABR').	'IBP'
sabr_args	dict	The arguments for the SABR method.	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 latent PGD search-range factor; also scales SABR input searches in STAPS. Nominal IBP regions are unchanged.	1.0

Returns:

Name	Type	Description
taps_bound	Tuple[Tensor, Tensor]	The TAPS bounds of the model's output.

Source code in CTRAIN/bound/taps.py

def bound_taps(original_model, hardened_model, bounded_blocks, data, target, n_classes, ptb, device='cuda', pgd_steps=20, pgd_restarts=1, pgd_step_size=.2,
               pgd_decay_factor=.2, pgd_decay_checkpoints=(5,7),
               gradient_link_thresh=.5, gradient_link_tolerance=1e-05, propagation="IBP", sabr_args=None, relu_upper_retention=1.0, pgd_eps_factor=1.0):
    """Estimate classifier margins using interval propagation and latent PGD.

    TAPS propagates the input region through the feature extractor using IBP,
    then searches for pivotal points in its output interval. STAPS uses SABR
    for feature propagation. The returned latent margins are training proxies;
    use ordinary full-region bounds for sound certification.

    Args:
        original_model (torch.nn.Module): The original neural network model.
        hardened_model (auto_LiRPA.BoundedModule): The auto_LiRPA model.
        bounded_blocks (list): List of bounded blocks of the model.
        data (Tensor): The input data tensor.
        target (Tensor): The target labels tensor.
        n_classes (int): The number of classes for classification.
        ptb (auto_LiRPA.PerturbationLpNorm): The perturbation object defining the
            perturbation set.
        device (str, optional): The device to run the computation on.
        pgd_steps (int, optional): The number of steps for the PGD attack.
        pgd_restarts (int, optional): The number of restarts for the PGD attack.
        pgd_step_size (float, optional): Latent-PGD step-size factor; None uses max(0.2, 2 /
            steps) for positive step counts.
        pgd_decay_factor (float, optional): The decay factor for the PGD attack.
        pgd_decay_checkpoints (tuple, optional): The decay checkpoints for the PGD attack.
        gradient_link_thresh (float, optional): The threshold for gradient linking.
        gradient_link_tolerance (float, optional): The tolerance for gradient linking.
        propagation (str, optional): The propagation method to use ('IBP' or 'SABR').
        sabr_args (dict, optional): The arguments for the SABR method.
        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 latent PGD search-range factor; also scales SABR
            input searches in STAPS. Nominal IBP regions are unchanged.

    Returns:
        taps_bound(Tuple[Tensor, Tensor]): The TAPS bounds of the model's output.
    """
    validate_hybrid_options(relu_upper_retention, pgd_eps_factor)
    with preserve_model_state(original_model, hardened_model, *bounded_blocks):
        assert len(bounded_blocks) == 2, "Split not supported!"

        if propagation == 'IBP':
            lb, ub = bound_ibp(
                model=bounded_blocks[0],
                ptb=ptb,
                data=data,
                target=None,
                n_classes=n_classes,
                bound_upper=True,
                relu_upper_retention=relu_upper_retention,
            )
        elif propagation == 'SABR':
            assert sabr_args is not None, "Need to Provide SABR arguments if you choose SABR for propagation"
            lb, ub = bound_sabr(
                # Intermediate Bound model instructs to return bounds after the first network block
                **{**sabr_args, "intermediate_bound_model": bounded_blocks[0], "return_adv_output": False,
                                  "relu_upper_retention": relu_upper_retention, "pgd_eps_factor": pgd_eps_factor},
            )

        else:
            raise ValueError("propagation must be IBP or SABR")

        c = construct_c(data, target, n_classes)
        # Attack in eval mode, then evaluate points in the caller's mode (De Palma et al.).
        with preserve_model_state(original_model, hardened_model, *bounded_blocks):
            for model in (original_model, hardened_model, *bounded_blocks):
                model.eval()
            with torch.no_grad():
                attack_lb, attack_ub = lb, ub
                if pgd_eps_factor != 1:
                    midpoint, radius = (lb + ub) / 2, (ub - lb) / 2
                    attack_lb, attack_ub = midpoint - radius * pgd_eps_factor, midpoint + radius * pgd_eps_factor
                adv_samples = _get_pivotal_points(bounded_blocks[1], attack_lb, attack_ub, pgd_steps, pgd_restarts, pgd_step_size, pgd_decay_factor, pgd_decay_checkpoints, n_classes, C=c)

        pts = adv_samples[0].detach()
        pts = torch.transpose(pts, 0, 1)
        pts = RectifiedLinearGradientLink.apply(lb.unsqueeze(0), ub.unsqueeze(0), pts, gradient_link_thresh, gradient_link_tolerance)
        pts = torch.transpose(pts, 0, 1)
        pgd_bounds = _get_bound_estimation_from_pts(bounded_blocks[1], pts, None, c)
        # NOTE: VERY IMPORTANT CHANGES TO TAPS BOUND TO BE COMPATIBLE WITH CTRAIN WORKFLOW
        pgd_bounds = pgd_bounds[:, 1:]
        pgd_bounds = -pgd_bounds


        ibp_lb, ibp_ub = bound_ibp(
            model=bounded_blocks[1],
            ptb=PerturbationLpNorm(x_L=lb, x_U=ub),
            data=(lb + ub) / 2,
            relu_upper_retention=relu_upper_retention,
            target=target,
            n_classes=n_classes,
        )

        return pgd_bounds, ibp_lb

split_bound_graph(graph, block_sizes, layers)

Expose feature-extractor and classifier views of one converted network.

The views share the original graph's operators, parameters and BatchNorm buffers. This avoids a second conversion when TAPS alternates between feature bounds and latent classifier searches.

Parameters:

Name	Type	Description	Default
graph	BoundedModule	Converted sequential single-input network.	required
block_sizes	tuple[int, int]	Positive layer counts covering every layer exactly once.	required
layers	Sequence	Sequential layers corresponding to the graph operators.	required

Returns:

Type	Description
list	Feature-extractor and classifier graph views.

Source code in CTRAIN/bound/taps.py

def split_bound_graph(graph, block_sizes, layers):
    """Expose feature-extractor and classifier views of one converted network.

    The views share the original graph's operators, parameters and BatchNorm
    buffers. This avoids a second conversion when TAPS alternates between feature
    bounds and latent classifier searches.

    Args:
        graph (auto_LiRPA.BoundedModule): Converted sequential single-input network.
        block_sizes (tuple[int, int]): Positive layer counts covering every layer exactly
            once.
        layers (Sequence): Sequential layers corresponding to the graph operators.

    Returns:
        (list): Feature-extractor and classifier graph views.
    """
    if (len(block_sizes) != 2
            or any(type(size) is not int or size <= 0 for size in block_sizes)
            or sum(block_sizes) != len(layers)):
        raise ValueError('TAPS/STAPS require two positive block sizes covering every layer')
    if sum(index is not None for index in graph.input_index) != 1:
        raise ValueError('TAPS/STAPS require a sequential single-input bound graph')
    chain, node = [], graph[graph.output_name[0]]
    while node.name != graph.input_name[0]:
        if not node.inputs or any(parent.from_input for parent in node.inputs[1:]):
            raise ValueError('TAPS/STAPS require a sequential single-input bound graph')
        chain.append(node)
        node = node.inputs[0]
    chain.reverse()
    if len(chain) != len(layers):
        raise ValueError('TAPS/STAPS require one bound operator per sequential layer')
    cut = chain[block_sizes[0] - 1]
    return [_BoundGraphBlock(graph, cut), _BoundGraphBlock(graph, chain[-1], cut)]