αβ-CROWN(神经网络形式化验证器)

概述

αβ-CROWN(读作”alpha-beta-CROWN”)是由斯坦福大学、清华大学和卡内基梅隆大学联合开发的高效神经网络验证器,专门用于神经网络对抗鲁棒性认证。它是CROWN验证器的增强版本,通过引入分支定界(Branch-and-Bound)策略和αβ参数化技术,实现了业界领先的性能。

核心特点:

  • 完全可验证:提供严格的鲁棒性下界
  • 高效:支持GPU加速
  • 可扩展:能够验证大规模网络
  • 开源:基于PyTorch

背景:神经网络验证问题

对抗鲁棒性

对于分类器 ,鲁棒性问题定义为:

即在以 为中心、 为半径的L_p球内,分类器输出应保持不变。

验证问题的数学形式

给定神经网络 、输入 和扰动范围 ,验证问题可以表述为:

如果最小值 > 0,则网络在 范围内是鲁棒的。

αβ-CROWN核心算法

1. CROWN界限传播

CROWN使用线性松弛来界定点经过神经网络后的输出:

import torch
import torch.nn as nn
from typing import Tuple, List
 
class CROWNBounds:
    """CROWN界限传播"""
    
    def __init__(self, network: nn.Module):
        self.network = network
        self.layers = list(network.children())
    
    def compute_bounds(self, x: torch.Tensor, eps: float, 
                      lower: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        计算输入扰动下的输出界限
        
        Args:
            x: 干净输入 [batch, dim]
            eps: 扰动半径
            lower: 是否计算下界
        
        Returns:
            (lower_bounds, upper_bounds)
        """
        lb = x - eps if lower else None
        ub = x + eps if not lower else None
        
        # 逐层传播界限
        for layer in self.layers:
            lb, ub = self._propagate_layer(layer, lb, ub)
        
        return lb, ub
    
    def _propagate_linear(self, layer: nn.Linear, 
                         lb: torch.Tensor, ub: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """线性层界限传播"""
        weight = layer.weight
        
        if weight.shape[0] == 1:  # 单输出
            w_pos = torch.clamp(weight, min=0)
            w_neg = torch.clamp(weight, max=0)
            
            lb_new = w_pos @ lb + w_neg @ ub
            ub_new = w_pos @ ub + w_neg @ lb
        else:
            # 多输出情况
            w_pos = torch.clamp(weight, min=0)
            w_neg = torch.clamp(weight, max=0)
            
            lb_new = w_pos @ lb.unsqueeze(1) + w_neg @ ub.unsqueeze(1)
            ub_new = w_pos @ ub.unsqueeze(1) + w_neg @ lb.unsqueeze(1)
        
        return lb_new, ub_new.squeeze(-1)
    
    def _propagate_relu(self, lb: torch.Tensor, ub: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """ReLU层界限传播"""
        # 简单的不确定性传播
        lb_pos = torch.clamp(lb, min=0)
        ub_neg = torch.clamp(ub, max=0)
        
        # 线性松弛
        slope = ub_pos / (ub_pos - lb_neg + 1e-8)
        bias = -slope * lb_neg
        
        lb_new = slope * lb + bias
        ub_new = torch.clamp(ub, min=0)
        
        return lb_new, ub_new

2. α参数化

α参数化是对CROWN界限的增强,通过引入可学习的α参数来收紧界限:

class AlphaCROWNBounds:
    """α-CROWN界限:可学习的α参数"""
    
    def __init__(self, network: nn.Module):
        self.network = network
        self.alpha = {}  # 每层ReLU的α参数
    
    def init_alpha(self, x: torch.Tensor, eps: float):
        """初始化α参数"""
        device = x.device
        layer_idx = 0
        
        with torch.no_grad():
            x_min = x - eps
            x_max = x + eps
            
            for name, module in self.network.named_modules():
                if isinstance(module, nn.ReLU):
                    # 初始化α为不确定性传播的值
                    self.alpha[layer_idx] = torch.zeros(1, device=device)
                    layer_idx += 1
    
    def compute_alpha_bounds(self, x: torch.Tensor, eps: float) -> Tuple:
        """计算α-CROWN界限"""
        device = x.device
        
        # 初始化输入范围
        lower = x - eps
        upper = x + eps
        
        biases = []
        slopes = []
        layer_idx = 0
        
        for module in self.network.modules():
            if isinstance(module, nn.Linear):
                weight = module.weight
                bias = module.bias
                
                # 线性变换界限
                w_pos = torch.clamp(weight, min=0)
                w_neg = torch.clamp(weight, max=0)
                
                lower = w_pos @ lower + w_neg @ upper
                upper = w_pos @ upper + w_neg @ lower
                
                if bias is not None:
                    lower = lower + bias
                    upper = upper + bias
                    
            elif isinstance(module, nn.ReLU):
                # ReLU的α参数化
                alpha = self.alpha.get(layer_idx, torch.zeros(1, device=device))
                
                # α-区间
                lower_alpha = alpha * lower
                upper_alpha = (1 - alpha) * upper
                
                # α-CROWN下界
                lb_new = torch.clamp(lower, min=0) + \
                        torch.clamp(lower_alpha, max=0) - \
                        torch.clamp(-lower_alpha, max=0)
                
                # α-CROWN上界
                ub_new = torch.clamp(upper, min=0) + \
                        torch.clamp(upper_alpha, min=0) - \
                        torch.clamp(-upper_alpha, max=0)
                
                biases.append((lower_alpha, upper_alpha))
                slopes.append(alpha)
                
                lower = lb_new
                upper = ub_new
                
                layer_idx += 1
        
        return lower, upper, biases, slopes

3. 分支定界(Branch and Bound)

当CROWN界限不足以证明鲁棒性时,αβ-CROWN使用分支定界策略:

class BranchAndBound:
    """分支定界算法用于神经网络验证"""
    
    def __init__(self, network: nn.Module, eps: float):
        self.network = network
        self.eps = eps
        self.bounds = AlphaCROWNBounds(network)
        
    def verify(self, x: torch.Tensor, y_true: int) -> Tuple[bool, float]:
        """
        验证网络鲁棒性
        
        Returns:
            (is_robust, certified_radius)
        """
        # 初始界限
        lb, ub, _, _ = self.bounds.compute_alpha_bounds(x, self.eps)
        
        # 检查是否存在对抗样本
        diff = ub - lb[y_true]
        
        if torch.all(diff <= 0):
            return True, self.eps
        
        # 分支定界
        queue = [(x, self.eps, lb, ub)]
        best_lower = 0.0
        
        while queue and len(queue) < 100000:  # 防止无限循环
            x_cur, eps_cur, lb_cur, ub_cur = queue.pop()
            
            # 找到最不确定的ReLU
            relu_idx = self._find_most_uncertain_relu(lb_cur, ub_cur)
            
            if relu_idx is None:
                continue
            
            # 分支:尝试两种情况
            eps1 = eps_cur / 2
            eps2 = eps_cur / 2
            
            # 分支1
            lb1, ub1, _, _ = self.bounds.compute_alpha_bounds(x_cur, eps1)
            if self._check_adversarial(lb1, ub1, y_true):
                continue  # 仍然存在对抗样本
            
            # 分支2
            lb2, ub2, _, _ = self.bounds.compute_alpha_bounds(x_cur, eps2)
            if self._check_adversarial(lb2, ub2, y_true):
                continue
            
            # 更新最优下界
            best_lower = max(best_lower, eps1)
            
            if best_lower >= self.eps:
                return True, self.eps
        
        return False, best_lower
    
    def _find_most_uncertain_relu(self, lb: torch.Tensor, 
                                   ub: torch.Tensor) -> int:
        """找到最不确定的ReLU"""
        # 计算不确定性区间
        uncertainty = ub - lb
        most_uncertain = torch.argmax(uncertainty)
        return most_uncertain.item()
    
    def _check_adversarial(self, lb: torch.Tensor, 
                          ub: torch.Tensor, y_true: int) -> bool:
        """检查是否存在对抗样本"""
        diff = ub - lb[y_true]
        return torch.any(diff > 0)

完整αβ-CROWN实现

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Dict, Optional
from dataclasses import dataclass
 
 
@dataclass
class LayerBounds:
    """每层的界限"""
    lower: torch.Tensor
    upper: torch.Tensor
    pre_lower: Optional[torch.Tensor] = None
    pre_upper: Optional[torch.Tensor] = None
 
 
class AlphaBetaCROWN:
    """
    完整的αβ-CROWN验证器
    """
    
    def __init__(self, model: nn.Module, device: str = 'cuda'):
        self.model = model
        self.device = device
        self.alpha = {}
        self.beta = {}
        
    def set_input_bounds(self, x: torch.Tensor, eps: float):
        """设置输入扰动范围"""
        self.x_center = x
        self.eps = eps
        self.input_lower = x - eps
        self.input_upper = x + eps
        
    def compute_bounds(self, stop_at_layer: int = None) -> Dict[int, LayerBounds]:
        """
        计算所有层的界限
        
        Args:
            stop_at_layer: 提前停止的层索引
        
        Returns:
            每层的界限字典
        """
        bounds = {}
        x_lower = self.input_lower
        x_upper = self.input_upper
        
        layer_idx = 0
        
        for name, module in self.model.named_modules():
            if stop_at_layer is not None and layer_idx >= stop_at_layer:
                break
                
            if isinstance(module, nn.Linear):
                w = module.weight.data
                b = module.bias.data if module.bias is not None else None
                
                # 线性层界限传播
                x_lower, x_upper = self._linear_bounds(
                    x_lower, x_upper, w, b
                )
                
            elif isinstance(module, nn.Conv2d):
                w = module.weight.data
                b = module.bias.data if module.bias is not None else None
                
                # 卷积层界限传播
                x_lower, x_upper = self._conv_bounds(
                    x_lower, x_upper, w, b, module.stride,
                    module.padding, module.dilation, module.groups
                )
                
            elif isinstance(module, nn.ReLU):
                # ReLU层:α参数化
                alpha = self.alpha.get(layer_idx, 
                    torch.zeros_like(x_lower).fill_(0.5))
                
                x_lower, x_upper, pre_lb, pre_ub = self._relu_bounds_alpha(
                    x_lower, x_upper, alpha
                )
                
                bounds[layer_idx] = LayerBounds(
                    lower=x_lower,
                    upper=x_upper,
                    pre_lower=pre_lb,
                    pre_upper=pre_ub
                )
                
                layer_idx += 1
                
            elif isinstance(module, nn.Flatten):
                x_lower = x_lower.view(x_lower.size(0), -1)
                x_upper = x_upper.view(x_upper.size(0), -1)
                
            elif isinstance(module, nn.BatchNorm2d):
                x_lower, x_upper = self._bn_bounds(
                    x_lower, x_upper, module
                )
        
        return bounds
    
    def _linear_bounds(self, lb: torch.Tensor, ub: torch.Tensor,
                       w: torch.Tensor, b: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """线性层界限传播"""
        w_pos = torch.clamp(w, min=0)
        w_neg = torch.clamp(w, max=0)
        
        if b is not None:
            new_lb = w_pos @ lb.unsqueeze(-1) + w_neg @ ub.unsqueeze(-1) + b.unsqueeze(-1)
            new_ub = w_pos @ ub.unsqueeze(-1) + w_neg @ lb.unsqueeze(-1) + b.unsqueeze(-1)
        else:
            new_lb = w_pos @ lb.unsqueeze(-1) + w_neg @ ub.unsqueeze(-1)
            new_ub = w_pos @ ub.unsqueeze(-1) + w_neg @ lb.unsqueeze(-1)
        
        return new_lb.squeeze(-1), new_ub.squeeze(-1)
    
    def _conv_bounds(self, lb: torch.Tensor, ub: torch.Tensor,
                    w: torch.Tensor, b: torch.Tensor,
                    stride: Tuple, padding: Tuple,
                    dilation: Tuple, groups: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """卷积层界限传播"""
        w_pos = torch.clamp(w, min=0)
        w_neg = torch.clamp(w, max=0)
        
        conv_pos = lambda x: F.conv2d(x, w_pos, b, stride, padding, dilation, groups)
        conv_neg = lambda x: F.conv2d(x, w_neg, b, stride, padding, dilation, groups)
        
        new_lb = conv_pos(lb) + conv_neg(ub)
        new_ub = conv_pos(ub) + conv_neg(lb)
        
        return new_lb, new_ub
    
    def _relu_bounds_alpha(self, lb: torch.Tensor, ub: torch.Tensor,
                          alpha: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """ReLU层的α参数化界限"""
        # 保存pre-activation界限
        pre_lb = lb.clone()
        pre_ub = ub.clone()
        
        # α-CROWN界限
        alpha = alpha.clamp(0, 1)
        
        # 下界
        lb_lower = torch.clamp(lb, min=0)
        lb_upper = torch.clamp(ub, min=0)
        
        # α区间
        lower_alpha = alpha * lb
        upper_alpha = (1 - alpha) * ub
        
        # 组合界限
        new_lb = lb_lower + torch.clamp(lower_alpha, max=0) - torch.clamp(-lower_alpha, max=0)
        new_ub = lb_upper + torch.clamp(upper_alpha, min=0) - torch.clamp(-upper_alpha, max=0)
        
        return new_lb, new_ub, pre_lb, pre_ub
    
    def get_lower_bound(self, target_class: int) -> float:
        """获取目标类别的鲁棒性下界"""
        bounds = self.compute_bounds()
        
        # 最后一层
        final_bounds = bounds[max(bounds.keys())]
        
        # 返回目标类别的下界
        return final_bounds.lower[:, target_class].min().item()
 
 
def verify_robustness(model: nn.Module, x: torch.Tensor, 
                       y_true: int, eps: float, device: str = 'cuda') -> Dict:
    """
    验证网络鲁棒性的便捷函数
    
    Returns:
        验证结果字典
    """
    model.eval()
    model.to(device)
    
    x = x.to(device).unsqueeze(0)
    
    verifier = AlphaBetaCROWN(model, device)
    verifier.set_input_bounds(x, eps)
    
    # 计算界限
    bounds = verifier.compute_bounds()
    
    # 提取输出界限
    final_layer_idx = max(bounds.keys())
    final_bounds = bounds[final_layer_idx]
    
    lb = final_bounds.lower[0]
    ub = final_bounds.upper[0]
    
    # 检查鲁棒性
    diff = ub - ub[y_true]
    diff[y_true] = float('inf')  # 忽略真实类别
    
    if torch.all(diff <= 0):
        return {
            'verified': True,
            'certified_radius': eps,
            'method': 'alpha-beta-crown'
        }
    else:
        # 尝试分支定界
        bab = BranchAndBound(model, eps)
        is_robust, certified_radius = bab.verify(x.squeeze(0), y_true)
        
        return {
            'verified': is_robust,
            'certified_radius': certified_radius,
            'method': 'branch-and-bound'
        }

使用示例

import torchvision.models as models
 
# 加载预训练模型
model = models.resnet18(pretrained=True)
model.eval()
 
# 准备输入
x = torch.randn(1, 3, 224, 224)
y_true = 0  # 真实类别
 
# 验证
eps = 0.01
result = verify_robustness(model, x, y_true, eps)
 
print(f"Verified: {result['verified']}")
print(f"Certified radius: {result['certified_radius']:.6f}")
print(f"Method: {result['method']}")

与其他验证器对比

验证器核心算法速度可扩展性精确度
αβ-CROWN线性松弛+分支定界非常快
ERAN抽象解释
MarabouSMT求解完全精确
POPAMGPU加速

应用场景

1. 自动驾驶安全认证

# 验证自动驾驶感知网络
def verify_autonomous_driving(model, sensor_input, eps=0.01):
    """验证感知网络对对抗扰动的鲁棒性"""
    predictions = model(sensor_input)
    top_class = predictions.argmax()
    
    result = verify_robustness(model, sensor_input, top_class.item(), eps)
    
    if result['verified']:
        print("感知网络通过安全认证")
    else:
        print(f"存在风险,可认证半径: {result['certified_radius']}")
    
    return result

2. 医疗AI验证

# 验证医疗诊断模型的鲁棒性
def verify_medical_ai(model, patient_data, eps=0.005):
    """验证医疗AI的鲁棒性"""
    result = verify_robustness(model, patient_data, 
                               model(patient_data).argmax().item(), 
                               eps)
    
    # 医疗场景需要更严格的认证
    required_radius = 0.01
    
    if result['verified'] or result['certified_radius'] >= required_radius:
        return "PASS"
    else:
        return "FAIL"

总结

αβ-CROWN是当前最高效的神经网络验证器之一,通过以下技术创新实现了优秀的性能:

  1. α参数化:可学习的界限参数提供更紧的松弛
  2. β参数化:进一步收紧界限定
  3. 分支定界:处理复杂网络
  4. GPU加速:利用并行计算

它为神经网络的安全部署提供了重要的理论保证。


参考资料