Set-based zonotope propagation API

CTRAIN.bound.set_based_zono contains the specialized representation and auto_LiRPA adapters used by set-based training. It retains centers, correlated generators and independent error radii for the training objective; it is not a general zonotope library.

See representation and operators for supported semantics and local LiRPA adapters.

Represent a correlated input or output enclosure for set-based training.

The represented points are c + sum_q G[:, q] * xi_q + r * eta, where each xi_q and each coordinate of eta lie in [-1, 1]. Shared generator coefficients preserve coordinate correlations, while r describes independent error.

Parameters:

Name	Type	Description	Default
c	Tensor	Centre of shape (batch, ...).	required
G	Tensor	Correlated generators of shape (batch, Q, ...).	required
r	Tensor	Nonnegative independent radii with the same shape as c.	required

Source code in CTRAIN/bound/set_based_zono.py

@dataclass
class Zonotope:
    """Represent a correlated input or output enclosure for set-based training.

    The represented points are c + sum_q G[:, q] * xi_q + r * eta, where each
    xi_q and each coordinate of eta lie in [-1, 1]. Shared generator coefficients
    preserve coordinate correlations, while r describes independent error.

    Args:
        c (torch.Tensor): Centre of shape (batch, ...).
        G (torch.Tensor): Correlated generators of shape (batch, Q, ...).
        r (torch.Tensor): Nonnegative independent radii with the same shape as c.
    """

    c: torch.Tensor
    G: torch.Tensor
    r: torch.Tensor

    def bounds(self):
        """Return coordinatewise lower and upper endpoints enclosing the zonotope."""
        radius = self.G.abs().sum(1) + self.r
        return self.c - radius, self.c + radius

bounds()

Return coordinatewise lower and upper endpoints enclosing the zonotope.

Source code in CTRAIN/bound/set_based_zono.py

def bounds(self):
    """Return coordinatewise lower and upper endpoints enclosing the zonotope."""
    radius = self.G.abs().sum(1) + self.r
    return self.c - radius, self.c + radius

Propagate a correlated enclosure through supported auto_LiRPA operators.

The returned Zonotope retains its centre, generators and independent error radii with an autograd graph for set-based training. Same-slope ReLU relaxations preserve shared generator coefficients. Nonlinear errors either remain independent intervals or become additional generators.

Parameters:

Name	Type	Description	Default
model	Module	Supported sequential network.	required
zonotope	Zonotope	Input enclosure in model coordinates.	required
error_propagation	str	'interval' or 'generators'.	'interval'
bn_statistics	dict	Differentiable clean means and variances indexed by BatchNorm layer. Otherwise evaluation uses running statistics.	None
layers	list	Previously validated sequential layers.	None

Returns:

Type	Description
Zonotope	Propagated output enclosure, retaining parameter gradients.

Source code in CTRAIN/bound/set_based_zono.py

def bound_zonotope(model, zonotope, error_propagation='interval', bn_statistics=None, layers=None):
    """Propagate a correlated enclosure through supported auto_LiRPA operators.

    The returned Zonotope retains its centre, generators and independent error
    radii with an autograd graph for set-based training. Same-slope ReLU
    relaxations preserve shared generator coefficients. Nonlinear errors either
    remain independent intervals or become additional generators.

    Args:
        model (torch.nn.Module): Supported sequential network.
        zonotope (Zonotope): Input enclosure in model coordinates.
        error_propagation (str): 'interval' or 'generators'.
        bn_statistics (dict, optional): Differentiable clean means and variances indexed by
            BatchNorm layer. Otherwise evaluation uses running
            statistics.
        layers (list, optional): Previously validated sequential layers.

    Returns:
        (Zonotope): Propagated output enclosure, retaining parameter gradients.
    """
    if error_propagation not in ('interval', 'generators'):
        raise ValueError('error_propagation must be interval or generators')
    layers = sequential_layers(model) if layers is None else layers
    c, G, r = zonotope.c, zonotope.G, zonotope.r
    if error_propagation == 'generators' and (r != 0).any():
        G = torch.cat((G, diagonal_generators(r)), 1)
        r = torch.zeros_like(r)
    linear = LinearBound(G, c - r, G, c + r)
    for layer in layers:
        linear = _forward_operator(layer, linear, bn_statistics)
        # Same-slope propagation keeps both coefficient tensors equal. Share them.
        linear.uw = linear.lw
        if error_propagation == 'generators' and isinstance(layer, nn.ReLU):
            # Installed dynamic-forward lacks BN support and ReLU scatter backward;
            # append errors locally so later layers retain them with working autograd.
            c, r = (linear.lb + linear.ub) / 2, (linear.ub - linear.lb) / 2
            G = linear.lw
            if (r != 0).any():
                G = torch.cat((G, diagonal_generators(r)), 1)
            linear = LinearBound(G, c, G, c)
    return Zonotope((linear.lb + linear.ub) / 2, linear.lw, (linear.ub - linear.lb) / 2)

Input boxes and output margins

box_zonotope builds this representation from finite, ordered, batched lower and upper endpoints, with one diagonal generator per input coordinate. zonotope_margins returns lower true-class-minus-competitor margins, retaining generator correlations. Its true-class entry is inf and can be ignored when checking whether every competing-class margin is positive.

For ordinary full-box classification certification, use the shared evaluation API with method="FORWARD".

Represent a finite, ordered input box using one generator per coordinate.

Returns a Zonotope centred at (lower + upper) / 2 with independent half-width generators and zero additional error radii.

Source code in CTRAIN/bound/set_based_zono.py

def box_zonotope(lower, upper):
    """Represent a finite, ordered input box using one generator per coordinate.

    Returns a Zonotope centred at (lower + upper) / 2 with independent half-width
    generators and zero additional error radii.
    """
    if lower.shape != upper.shape or not torch.isfinite(lower).all() or not torch.isfinite(upper).all() or (lower > upper).any():
        raise ValueError('Box endpoints must be finite, ordered, and have matching shapes')
    return Zonotope((lower + upper) / 2, diagonal_generators((upper - lower) / 2),
                    torch.zeros_like(lower))

Compute lower classification margins while retaining generator correlations.

Shared generator contributions may cancel when subtracting competing logits; independent error radii are added. Returns one column per class, with the true-class column set to infinity.

Source code in CTRAIN/bound/set_based_zono.py

def zonotope_margins(zonotope, labels):
    """Compute lower classification margins while retaining generator correlations.

    Shared generator contributions may cancel when subtracting competing logits;
    independent error radii are added. Returns one column per class, with the
    true-class column set to infinity.
    """
    c, G, r = zonotope.c, zonotope.G, zonotope.r
    if c.ndim != 2:
        raise ValueError('Certification requires a batch of class logits')
    index = labels[:, None]
    gy = G.gather(2, index[:, None].expand(-1, G.shape[1], 1))
    # Shared generator contributions can cancel; independent radii must be added.
    margins = c.gather(1, index) - c - (gy - G).abs().sum(1) - r.gather(1, index) - r
    return margins.masked_fill(F.one_hot(labels, c.shape[1]).bool(), float('inf'))