configspace_abcrown.py

from ConfigSpace import ConfigurationSpace, EqualsCondition, Float, ForbiddenAndConjunction, ForbiddenEqualsClause, ForbiddenInClause, InCondition, Integer, Categorical, Constant, OrdinalHyperparameter
import torch

def build_abcrown_config_space(eps, no_classes, include_bab=True, include_mip=True,
                               include_bab_refine=True, include_input_split=True,
                               include_attack=True, tune_attack_extensive=False):

    if isinstance(eps, torch.Tensor):
        eps = eps.cpu().numpy()

    abcrown_cs = ConfigurationSpace()

    complete_verifiers = []
    if include_bab:
        complete_verifiers.append('bab')
    if include_bab_refine:
        complete_verifiers.append('bab-refine')
    if include_mip:
        complete_verifiers.append('mip')

    # ==================================
    # Space Definitions
    # ==================================

    # 1. General Space
    abcrown_general_cs = ConfigurationSpace(
        name='general',
        space={
            "device": Constant(name='device', value="cuda") if torch.cuda.is_available() else Constant(name='device', value='cpu'),
            "conv_mode": Categorical("conv_mode", ["patches", "matrix"], default="patches"),
            "complete_verifier": Categorical("complete_verifier", complete_verifiers, default=complete_verifiers[0]), 
            "enable_incomplete_verification": [True, False], 
        }
    )
    abcrown_cs.add_configuration_space(prefix='general', configuration_space=abcrown_general_cs)

    # 2. Solver Space
    abcrown_solver_cs = ConfigurationSpace(
        name='solver',
        space={
            "batch_size": OrdinalHyperparameter("batch_size", [1, 16, 64, 256, 1024, 4096, 16384, 65536], default_value=64),
            "bound_prop_method": Categorical("bound_prop_method", ["alpha-crown", "crown", "forward+crown"], default="alpha-crown"),
        }
    )
    abcrown_cs.add_configuration_space(prefix='solver', configuration_space=abcrown_solver_cs)

    # Note: solver:crown space is empty in the original file
    abcrown_crown_cs = ConfigurationSpace(name="crown", space={})
    abcrown_cs.add_configuration_space(prefix='solver:crown', configuration_space=abcrown_crown_cs)

    # 3. Alpha Crown Space
    abcrown_alpha_crown_cs = ConfigurationSpace(
        name='alpha_crown',
        space={
            "alpha": [True],
            "lr_alpha": Float("lr_alpha", (0.001, .5), log=True, default=.1),
            "iteration": Integer("iteration", (1, 500), default=100, log=True),
            "share_alphas": [False, True],
            "lr_decay": Float("lr_decay", (0.8, 1), default=.98, log=True),
            "full_conv_alpha": [True, False],
            "matmul_share_alphas": [False],
        }
    )
    abcrown_cs.add_configuration_space(prefix='solver:alpha-crown', configuration_space=abcrown_alpha_crown_cs)

    # 4. Beta Crown Space
    abcrown_betacrown_cs = ConfigurationSpace(
        name="beta_crown",
        space={
            "lr_alpha": Float("lr_alpha", (0.001, .25), log=True, default=.01),
            "lr_beta": Float("lr_beta", (0.001, 0.5), log=True, default=.05),
            "lr_decay": Float("lr_decay", (0.8, 1), default=.98, log=True),
            "iteration": Integer("iteration", (1, 100), default=50, log=True),
            "enable_opt_interm_bounds": [False],
            "all_node_split_LP": [False],
        }
    )
    abcrown_cs.add_configuration_space(prefix='solver:beta-crown', configuration_space=abcrown_betacrown_cs)

    # 5. MIP Space
    abcrown_mip_cs = ConfigurationSpace(
        name='mip',
        space={
            "refine_neuron_timeout": Integer("refine_neuron_timeout", (1, 30), default=15),
            "refine_neuron_time_percentage": Float("refine_neuron_time_percentage", (0.1, 0.9), default=.8),
        }
    )
    if include_bab_refine:
        abcrown_cs.add_configuration_space(prefix='solver:mip', configuration_space=abcrown_mip_cs)

    # 6. BaB Space
    abcrown_bab_cs = ConfigurationSpace(
        name="bab",
        space={
            "pruning_in_iteration": [True, False],
            "pruning_in_iteration_ratio": Float("pruning_in_iteration_ratio", (0, 1), default=.2),
            "sort_targets": [False],
            "batched_domain_list": [True],
            "vanilla_crown": [False],
            "interm_transfer": [True, False],
            'tree_traversal': Categorical('tree_traversal', ['breadth_first', 'depth_first'], default='depth_first'),
        }
    )
    if include_bab or include_bab_refine:
        abcrown_cs.add_configuration_space(prefix='bab', configuration_space=abcrown_bab_cs)

    # 7. BaB Cuts
    abcrown_cuts_cs = ConfigurationSpace(
        name='cut',
        space={
            'enabled': [False, True],
            'bab_cut': [True],
            'cplex_cuts': [False],
        }
    )
    abcrown_biccos_cs = ConfigurationSpace(
        name='biccos',
        space={
            'enabled': [True],
            # 'recursively_strengthening': [False, True],
            'drop_ratio': Float('drop_ratio', (0.0, 1.0), default=0.5),
            # 'verified_bonus': Float('verified_bonus', (0.0, 1.0), default=0.3),
            'max_infer_iter': Integer('max_infer_iter', (1, 100), default=20),
        }
    )
    abcrown_biccos_multi_tree_branching_cs = ConfigurationSpace(
        name='multi_tree_branching',
        space={
            'enabled': [True, False],
            'restore_best_tree': [True],
            # 'k_splits': Integer('k_splits', (1, 10), default=1),
            # 'keep_n_best_domains': Integer('keep_n_best_domains', (1, 100), default=50),
            # 'target_batch_size': Integer('target_batch_size', (1, 4000), default=200, log=False),
            'iterations': Integer('iterations', (1, 10), default=3),
        }
    )
    
    if include_bab:
        abcrown_cs.add_configuration_space(prefix='bab:cut', configuration_space=abcrown_cuts_cs)
        abcrown_cs.add_configuration_space(prefix='bab:cut:biccos', configuration_space=abcrown_biccos_cs)
        abcrown_cs.add_configuration_space(prefix='bab:cut:biccos:multi_tree_branching', configuration_space=abcrown_biccos_multi_tree_branching_cs)

    # 8. BaB Branching
    abcrown_branching_cs = ConfigurationSpace(
        name='branching',
        space={
            # "method": Categorical('method', ['kfsb', 'fsb', 'sb', 'kfsb-intercept-only', 'babsr', 'naive', 'brute-force'], default='kfsb'),
            "method": Categorical('method', ['kfsb', 'fsb', 'sb', 'kfsb-intercept-only', 'babsr', 'naive'], default='kfsb'),
            "candidates": Integer("candidates", (1, 20), default=3), # TODO: this should only be active when branching heuristic is fsb or kfsb
            "reduceop": ["min", "max"],
            "enable_intermediate_bound_opt": [False],
        }
    )
    if include_bab or include_bab_refine:
        abcrown_cs.add_configuration_space(prefix='bab:branching', configuration_space=abcrown_branching_cs)

    # 9. BaB Input Split
    abcrown_bab_input_split_cs = ConfigurationSpace(
        name='input_split',
        space={
            "enable": [False, True],
            # "enhanced_bound_prop_method": ["alpha-crown"],
            # "enhanced_branching_method": ["naive"],
            "adv_check": Categorical("adv_check", [-1, 0], default=0),
            "sb_sum": [False],
            "bf_zero_crossing_score": [False],
            "bf_iters": Constant("bf_iters", 1000000000),   
            "bf_batch_size": Constant("bf_batch_size", 100000),
            'ibp_enhancement': [False],
            "update_rhs_with_attack": [False],
            "sb_coeff_thresh": Float("sb_coeff_thresh", (0.0001, 2.0), log=True, default=.001),
            "sort_descending": [True], 
        }
    )
    if include_input_split:
        abcrown_cs.add_configuration_space(prefix='bab:branching:input_split', configuration_space=abcrown_bab_input_split_cs)

    # 9.5. BaB Clip N Verify
    abcrown_clip_n_verify_cs = ConfigurationSpace(
        name='clip_n_verify',
        space={
            # "rearrange_constraints": [False, True],
            # "alpha_crown": [False, True],
        }
    )
    abcrown_clip_input_domain_cs = ConfigurationSpace(
        name='clip_input_domain',
        space={
            "enabled": [False, True],
            "clip_type": Categorical("clip_type", ["complete", "relaxed"], default="relaxed"),
            # "clip_neuron_selection_type": Categorical("clip_neuron_selection_type", ["ratio", "number"], default="ratio"),
            # "clip_neuron_selection_value": Float("clip_neuron_selection_value", (0.0, 1.0), default=0.1),
            # "clip_iterations": Integer("clip_iterations", (1, 5), default=1),
        }
    )
    abcrown_clip_interm_domain_cs = ConfigurationSpace(
        name='clip_interm_domain',
        space={
            "enabled": [False, True],
            # "with_input": [False, True],
            "clip_in_alpha_crown": [False, True],
            # "topk_objective": Integer("topk_objective", (1, 100), default=20),
        }
    )
    if include_bab:
        abcrown_cs.add_configuration_space(prefix='bab:clip_n_verify', configuration_space=abcrown_clip_n_verify_cs)
        abcrown_cs.add_configuration_space(prefix='bab:clip_n_verify:clip_interm_domain', configuration_space=abcrown_clip_interm_domain_cs)
        if include_input_split:
            abcrown_cs.add_configuration_space(prefix='bab:clip_n_verify:clip_input_domain', configuration_space=abcrown_clip_input_domain_cs)

    # 10. Attack Space
    # construct differently depending on whether we are tuning extensively
    if include_attack:
        if tune_attack_extensive:
            abcrown_attack_cs = ConfigurationSpace(
                name="attack",
                space={
                    "pgd_steps": Integer("pgd_steps", (1, 1000), default=100, log=True),
                    "pgd_restarts": Integer("pgd_restarts", (1, 10_000), default=30, log=True),
                    "pgd_lr_decay": Float("pgd_lr_decay", (0.5, 0.99), default=.99, log=True),
                    # range depends on eps; this will be overridden when presets
                    # are used, but here it's valid when sampling freely.
                    "pgd_alpha": Float("pgd_alpha", (0.001 * eps, 2*eps), default=eps/4, log=True),
                    "pgd_alpha_scale": [False],
                    "enable_mip_attack": [False, True],
                    "attack_mode": ['PGD', 'diversed_PGD'],
                    "pgd_order": ['before', 'after'],
                }
            )
        else:
            # only expose a coarse-grained preset selector
            abcrown_attack_cs = ConfigurationSpace(
                name="attack",
                space={
                    "attack_preset": OrdinalHyperparameter("attack_preset",
                                                 ["weak", "strong", "super_strong", "extreme"],
                                                 default_value="strong"),
                    'pgd_order': Categorical('pgd_order', ['before', 'after'], default='before'),
                }
            )
        abcrown_cs.add_configuration_space(prefix='attack', configuration_space=abcrown_attack_cs)


    # ==================================
    # Conditions
    # ==================================
    overall_conditions = []
    
    # Attack conditions
    if include_attack:
        # no special conditions are required because the configuration space is
        # built differently depending on the ``tune_attack_extensive`` flag; the
        # unwanted variables are simply never present.
        pass
        
    # Input Split conditions
    if include_input_split:
        overall_conditions.extend([
            InCondition(abcrown_cs['bab:branching:input_split:enable'], abcrown_cs['general:complete_verifier'], [x for x in complete_verifiers if x != 'mip']),
            # EqualsCondition(abcrown_cs['bab:branching:input_split:enhanced_bound_prop_method'], abcrown_cs['bab:branching:input_split:enable'], True),
            # EqualsCondition(abcrown_cs['bab:branching:input_split:enhanced_branching_method'], abcrown_cs['bab:branching:input_split:enable'], True),
            EqualsCondition(abcrown_cs['bab:branching:input_split:adv_check'], abcrown_cs['bab:branching:input_split:enable'], True),
            
            EqualsCondition(abcrown_cs['bab:branching:input_split:sb_sum'], abcrown_cs['bab:branching:method'], "sb"),
            EqualsCondition(abcrown_cs['bab:branching:input_split:sb_coeff_thresh'], abcrown_cs['bab:branching:method'], "sb"),
            
            EqualsCondition(abcrown_cs['bab:branching:input_split:bf_zero_crossing_score'], abcrown_cs['bab:branching:method'], "naive"),
            EqualsCondition(abcrown_cs['bab:branching:input_split:bf_iters'], abcrown_cs['bab:branching:method'], "naive"),
            EqualsCondition(abcrown_cs['bab:branching:input_split:bf_batch_size'], abcrown_cs['bab:branching:method'], "naive"),
            
            EqualsCondition(abcrown_cs['bab:branching:input_split:ibp_enhancement'], abcrown_cs['bab:branching:input_split:enable'], True),
            EqualsCondition(abcrown_cs['bab:branching:input_split:update_rhs_with_attack'], abcrown_cs['bab:branching:input_split:enable'], True),
            EqualsCondition(abcrown_cs['bab:branching:input_split:sort_descending'], abcrown_cs['bab:branching:input_split:enable'], True),
        ])

    # BaB conditions
    if include_bab:
        for param in abcrown_cs.keys():
            if 'bab:' in param and 'bab:attack' not in param and 'bab:branching:input_split' not in param and 'bab:cut' not in param and 'bab:clip_n_verify' not in param and param not in ['bab:pruning_in_iteration_ratio', 'bab:sort_domain_interval']:
                overall_conditions.append(
                    InCondition(abcrown_cs[param], abcrown_cs['general:complete_verifier'], [x for x in complete_verifiers if x != 'mip'])
                )
        
        overall_conditions.extend([
            EqualsCondition(abcrown_cs['bab:pruning_in_iteration_ratio'], abcrown_cs['bab:pruning_in_iteration'], True),
        ])
        
        for param in abcrown_cs.keys():
            if 'bab:cut:biccos' in param:
                overall_conditions.append(EqualsCondition(abcrown_cs[param], abcrown_cs['bab:cut:enabled'], True))
                
        overall_conditions.append(EqualsCondition(abcrown_cs['bab:cut:enabled'], abcrown_cs['general:complete_verifier'], 'bab'))
    
    # Clip n Verify conditions
    if include_bab:
        if include_input_split:
            overall_conditions.append(InCondition(
                abcrown_cs['bab:clip_n_verify:clip_input_domain:enabled'], abcrown_cs['general:complete_verifier'], [x for x in complete_verifiers if x != 'mip']
            ))
        
        overall_conditions.append(
            InCondition(
                abcrown_cs['bab:clip_n_verify:clip_interm_domain:enabled'], abcrown_cs['general:complete_verifier'], [x for x in complete_verifiers if x != 'mip']
            )
        )
        
        overall_conditions.extend([
            EqualsCondition(
                abcrown_cs[param], abcrown_cs['bab:clip_n_verify:clip_interm_domain:enabled'], True
            ) for param in abcrown_cs.keys() if 'bab:clip_n_verify:clip_interm_domain:' in param and param != 'bab:clip_n_verify:clip_interm_domain:enabled'
        ])
        
        if include_input_split:
            overall_conditions.extend([
                EqualsCondition(
                    abcrown_cs[param], abcrown_cs['bab:clip_n_verify:clip_input_domain:enabled'], True
                ) for param in abcrown_cs.keys() if 'bab:clip_n_verify:clip_input_domain:' in param and param != 'bab:clip_n_verify:clip_input_domain:enabled'
            ])
        
    # Solver / Bound Prop / Beta Crown conditions
    if include_bab or include_bab_refine:
        overall_conditions.append(InCondition(abcrown_cs['solver:batch_size'], abcrown_cs['general:complete_verifier'], [x for x in complete_verifiers if x != 'mip']))
        
        for param in abcrown_cs.keys():
            if 'solver:alpha-crown:' in param:
                overall_conditions.append(EqualsCondition(abcrown_cs[param], abcrown_cs['solver:bound_prop_method'], 'alpha-crown'))
            elif 'solver:beta-crown' in param:
                overall_conditions.append(InCondition(abcrown_cs[param], abcrown_cs['general:complete_verifier'], [x for x in complete_verifiers if x != 'mip']))
    
    # BaB-refine MIP conditions
    if include_bab_refine:
        overall_conditions.extend([
            EqualsCondition(abcrown_cs['solver:mip:refine_neuron_time_percentage'], abcrown_cs['general:complete_verifier'], 'bab-refine'),
            EqualsCondition(abcrown_cs['solver:mip:refine_neuron_timeout'], abcrown_cs['general:complete_verifier'], 'bab-refine'),
        ])

    # General conditions
    overall_conditions.append(InCondition(abcrown_cs['general:conv_mode'], abcrown_cs['general:complete_verifier'], [x for x in complete_verifiers if x != 'mip']))

    abcrown_cs.add_conditions(overall_conditions)

    # ==================================
    # Forbidden Clauses
    # ==================================
    overall_forbidden = []

    if include_input_split:
        input_split_forbidden = [
            ForbiddenAndConjunction(
                ForbiddenEqualsClause(abcrown_cs['bab:branching:input_split:enable'], True),
                ForbiddenInClause(abcrown_cs['bab:branching:method'], ['kfsb', "babsr", "fsb", "kfsb-intercept-only"])
            ),
            ForbiddenAndConjunction(
                ForbiddenInClause(abcrown_cs['bab:branching:method'], ['sb', 'naive']),
                ForbiddenEqualsClause(abcrown_cs['bab:branching:input_split:enable'], False)
            ),
            ForbiddenAndConjunction(
                ForbiddenEqualsClause(abcrown_cs['bab:branching:input_split:enable'], False),
                ForbiddenEqualsClause(abcrown_cs['bab:clip_n_verify:clip_input_domain:enabled'], True)
            ),
            ForbiddenAndConjunction(
                ForbiddenEqualsClause(abcrown_cs['bab:branching:input_split:enable'], True),
                ForbiddenEqualsClause(abcrown_cs['bab:clip_n_verify:clip_interm_domain:enabled'], True)
            ),
        ]
        if 'mip' in complete_verifiers or 'bab-refine' in complete_verifiers:
            input_split_forbidden.append(
                ForbiddenAndConjunction(
                    ForbiddenEqualsClause(abcrown_cs['bab:branching:input_split:enable'], True),
                    ForbiddenInClause(abcrown_cs['general:complete_verifier'], [x for x in complete_verifiers if x != 'bab'])  
                )
            )
        
        overall_forbidden.extend(input_split_forbidden)
    
    if 'bab' in complete_verifiers and ('mip' in complete_verifiers or 'bab-refine' in complete_verifiers):
        overall_forbidden.extend([
            ForbiddenAndConjunction(
                ForbiddenEqualsClause(abcrown_cs['bab:interm_transfer'], False),
                ForbiddenInClause(abcrown_cs['general:complete_verifier'], [x for x in complete_verifiers if x != 'bab'])
            ),
            ForbiddenAndConjunction(
                ForbiddenEqualsClause(abcrown_cs['general:enable_incomplete_verification'], False),
                ForbiddenInClause(abcrown_cs['general:complete_verifier'], [x for x in complete_verifiers if x != 'bab'])   
            )
        ])
    
    if 'bab' in complete_verifiers:
        overall_forbidden.append(
            ForbiddenAndConjunction(
                ForbiddenEqualsClause(abcrown_cs['bab:cut:enabled'], True),
                ForbiddenInClause(abcrown_cs['bab:branching:method'], ['fsb', 'sb', 'babsr', 'naive'])
            )
        )
        if not include_input_split:
            overall_forbidden.append(
                ForbiddenInClause(abcrown_cs['bab:branching:method'], ['sb', 'naive'])
            )        

    abcrown_cs.add_forbidden_clauses(overall_forbidden)
    
    return abcrown_cs


if __name__ == "__main__":
    cs = build_abcrown_config_space(eps=.3, no_classes=10)
    print(cs.sample_configuration())