Implement BaBSR Branching Heuristic as a new Branching Strategy

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Next Release
   - Adding incremental infrastructure which allows pushing and popping constraints to/from the InputQuery.
   - Dropped support for parsing Tensorflow network format. Newest Marabou version that supports Tensorflow is at commit 190555573e4702.
+  - Added the BaBSR heuristic as a new branching strategy for ReLU Splitting
 
 ## Version 2.0.0
 

src/configuration/GlobalConfiguration.cpp

@@ -106,6 +106,7 @@ const unsigned GlobalConfiguration::REFACTORIZATION_THRESHOLD = 100;
 const GlobalConfiguration::BasisFactorizationType GlobalConfiguration::BASIS_FACTORIZATION_TYPE =
     GlobalConfiguration::SPARSE_FORREST_TOMLIN_FACTORIZATION;
 
+const unsigned GlobalConfiguration::BABSR_CANDIDATES_THRESHOLD = 5;
 const unsigned GlobalConfiguration::POLARITY_CANDIDATES_THRESHOLD = 5;
 
 const unsigned GlobalConfiguration::DNC_DEPTH_THRESHOLD = 5;

src/configuration/GlobalConfiguration.h

@@ -228,6 +228,11 @@ class GlobalConfiguration
     };
     static const BasisFactorizationType BASIS_FACTORIZATION_TYPE;
 
+    /* In the BaBSR-based branching heuristics, only this many earliest nodes are considered to
+       branch on.
+    */
+    static const unsigned BABSR_CANDIDATES_THRESHOLD;
+
     /* In the polarity-based branching heuristics, only this many earliest nodes
        are considered to branch on.
     */

src/configuration/OptionParser.cpp

@@ -152,7 +152,7 @@ void OptionParser::initialize()
             &( ( *_stringOptions )[Options::SPLITTING_STRATEGY] ) )
             ->default_value( ( *_stringOptions )[Options::SPLITTING_STRATEGY] ),
         "The branching strategy "
-        "(earliest-relu/pseudo-impact/largest-interval/relu-violation/polarity)."
+        "(earliest-relu/pseudo-impact/largest-interval/relu-violation/polarity/babsr-heuristic)."
         " pseudo-impact is specific to the DeepSoI (default) procedure and relu-violation is "
         "specific to the Reluplex procedure.\n" )(
         "soi-split-threshold",

src/configuration/Options.cpp

@@ -165,6 +165,8 @@ DivideStrategy Options::getDivideStrategy() const
         return DivideStrategy::LargestInterval;
     else if ( strategyString == "pseudo-impact" )
         return DivideStrategy::PseudoImpact;
+    else if ( strategyString == "babsr-heuristic" )
+        return DivideStrategy::BaBsrHeuristic;
     else
         return DivideStrategy::Auto;
 }

src/engine/CMakeLists.txt

@@ -47,6 +47,7 @@ engine_add_unit_test(SoftmaxConstraint)
 engine_add_unit_test(SmtCore)
 engine_add_unit_test(SumOfInfeasibilitiesManager)
 engine_add_unit_test(Tableau)
+engine_add_unit_test(BaBsrSplitting)
 
 if (${BUILD_PYTHON})
     target_include_directories(${MARABOU_PY} PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}")

src/engine/DivideStrategy.h

@@ -19,12 +19,13 @@
 enum class DivideStrategy {
     // Relu splitting
     Polarity = 0,    // Pick the ReLU with the polarity closest to 0 among the first K nodes
+    BaBsrHeuristic,  // Pick the ReLU with the highest BaBSR score
     EarliestReLU,    // Pick a ReLU that appears in the earliest layer
     ReLUViolation,   // Pick the ReLU that has been violated for the most times
     LargestInterval, // Pick the largest interval every K split steps, use ReLUViolation in other
                      // steps
     PseudoImpact,    // The pseudo-impact heuristic associated with SoI.
-    Auto,            // See decideBranchingHeursitics() in Engine.h
+    Auto,            // See decideBranchingHeuristics() in Engine.h
 };
 
 #endif // __DivideStrategy_h__

src/engine/Engine.cpp

@@ -2704,6 +2704,77 @@ void Engine::decideBranchingHeuristics()
     _smtCore.initializeScoreTrackerIfNeeded( _plConstraints );
 }
 
+PiecewiseLinearConstraint *Engine::pickSplitPLConstraintBasedOnBaBsrHeuristic()
+{
+    ENGINE_LOG( Stringf( "Using BaBsr heuristic with normalized scores..." ).ascii() );
+
+    if ( !_networkLevelReasoner )
+        return NULL;
+
+    // Get constraints from NLR
+    List<PiecewiseLinearConstraint *> constraints =
+        _networkLevelReasoner->getConstraintsInTopologicalOrder();
+
+    // Temporary vector to store raw scores for normalization
+    Vector<double> rawScores;
+    Map<double, PiecewiseLinearConstraint *> scoreToConstraint;
+
+    // Filter for ReLU constraints
+    for ( auto &plConstraint : constraints )
+    {
+        if ( plConstraint->supportBaBsr() && plConstraint->isActive() &&
+             !plConstraint->phaseFixed() )
+        {
+            ReluConstraint *reluConstraint = dynamic_cast<ReluConstraint *>( plConstraint );
+            if ( reluConstraint )
+            {
+                // Set NLR if not already set
+                reluConstraint->setNetworkLevelReasoner( _networkLevelReasoner );
+
+                // Calculate heuristic score - bias calculation happens inside computeBaBsr
+                plConstraint->updateScoreBasedOnBaBsr();
+
+                // Collect raw scores
+                rawScores.append( plConstraint->getScore() );
+            }
+        }
+    }
+
+    // Normalize scores
+    if ( !rawScores.empty() )
+    {
+        double minScore = *std::min_element( rawScores.begin(), rawScores.end() );
+        double maxScore = *std::max_element( rawScores.begin(), rawScores.end() );
+        double range = maxScore - minScore;
+
+        for ( auto &plConstraint : constraints )
+        {
+            if ( plConstraint->supportBaBsr() && plConstraint->isActive() &&
+                 !plConstraint->phaseFixed() )
+            {
+                double rawScore = plConstraint->getScore();
+                double normalizedScore = range > 0 ? ( rawScore - minScore ) / range : 0;
+
+                // Store normalized score in the map
+                scoreToConstraint[normalizedScore] = plConstraint;
+                if ( scoreToConstraint.size() >= GlobalConfiguration::BABSR_CANDIDATES_THRESHOLD )
+                    break;
+            }
+        }
+    }
+
+    // Split on neuron or constraint with highest normalized score
+    if ( !scoreToConstraint.empty() )
+    {
+        ENGINE_LOG( Stringf( "Normalized score of the picked ReLU: %f",
+                             ( *scoreToConstraint.begin() ).first )
+                        .ascii() );
+        return ( *scoreToConstraint.begin() ).second;
+    }
+    else
+        return NULL;
+}
+
 PiecewiseLinearConstraint *Engine::pickSplitPLConstraintBasedOnPolarity()
 {
     ENGINE_LOG( Stringf( "Using Polarity-based heuristics..." ).ascii() );
@@ -2815,6 +2886,8 @@ PiecewiseLinearConstraint *Engine::pickSplitPLConstraint( DivideStrategy strateg
                 candidatePLConstraint = _smtCore.getConstraintsWithHighestScore();
         }
     }
+    else if ( strategy == DivideStrategy::BaBsrHeuristic )
+        candidatePLConstraint = pickSplitPLConstraintBasedOnBaBsrHeuristic();
     else if ( strategy == DivideStrategy::Polarity )
         candidatePLConstraint = pickSplitPLConstraintBasedOnPolarity();
     else if ( strategy == DivideStrategy::EarliestReLU )
@@ -2837,7 +2910,9 @@ PiecewiseLinearConstraint *Engine::pickSplitPLConstraint( DivideStrategy strateg
 PiecewiseLinearConstraint *Engine::pickSplitPLConstraintSnC( SnCDivideStrategy strategy )
 {
     PiecewiseLinearConstraint *candidatePLConstraint = NULL;
-    if ( strategy == SnCDivideStrategy::Polarity )
+    if ( strategy == SnCDivideStrategy::BaBsrHeuristic )
+        candidatePLConstraint = pickSplitPLConstraintBasedOnBaBsrHeuristic();
+    else if ( strategy == SnCDivideStrategy::Polarity )
         candidatePLConstraint = pickSplitPLConstraintBasedOnPolarity();
     else if ( strategy == SnCDivideStrategy::EarliestReLU )
         candidatePLConstraint = pickSplitPLConstraintBasedOnTopology();

src/engine/Engine.h

@@ -751,6 +751,11 @@ class Engine
     */
     void decideBranchingHeuristics();
 
+    /*
+      Pick the ReLU with the highest BaBSR heuristic score.
+    */
+    PiecewiseLinearConstraint *pickSplitPLConstraintBasedOnBaBsrHeuristic();
+
     /*
       Among the earliest K ReLUs, pick the one with Polarity closest to 0.
       K is equal to GlobalConfiguration::POLARITY_CANDIDATES_THRESHOLD

src/engine/PiecewiseLinearConstraint.h

@@ -315,6 +315,14 @@ class PiecewiseLinearConstraint : public ITableau::VariableWatcher
         return false;
     }
 
+    /*
+      Return true if and only if this piecewise linear constraint supports the BaBsr Heuristic
+    */
+    virtual bool supportBaBsr() const
+    {
+        return false;
+    }
+
     /*
       Update the preferred direction to take first when splitting on this PLConstraint
     */
@@ -327,6 +335,9 @@ class PiecewiseLinearConstraint : public ITableau::VariableWatcher
         return _score;
     }
 
+    virtual void updateScoreBasedOnBaBsr()
+    {
+    }
 
     virtual void updateScoreBasedOnPolarity()
     {

src/engine/ReluConstraint.cpp

@@ -1019,6 +1019,11 @@ bool ReluConstraint::supportPolarity() const
     return true;
 }
 
+bool ReluConstraint::supportBaBsr() const
+{
+    return true;
+}
+
 bool ReluConstraint::auxVariableInUse() const
 {
     return _auxVarInUse;
@@ -1029,6 +1034,57 @@ unsigned ReluConstraint::getAux() const
     return _aux;
 }
 
+std::unordered_map<const ReluConstraint *, double> ReluConstraint::_biasCache;
+
+void ReluConstraint::initializeBiasCache( NLR::NetworkLevelReasoner &nlr )
+{
+    // Loop through all constraints in topological order
+    for ( const auto *constraint : nlr.getConstraintsInTopologicalOrder() )
+    {
+        // Only handle ReluConstraints and cache their biases
+        if ( const auto *reluConstraint = dynamic_cast<const ReluConstraint *>( constraint ) )
+        {
+            if ( _biasCache.find( reluConstraint ) == _biasCache.end() )
+            {
+                // Compute the bias and store it in the cache
+                double bias = nlr.getPrevBiasForReluConstraint( reluConstraint );
+                _biasCache[reluConstraint] = bias;
+            }
+        }
+    }
+}
+
+
+double ReluConstraint::computeBaBsr() const
+{
+    double biasTerm = calculateBias();
+
+    // get upper and lower bounds
+    double ub = getUpperBound( _b );
+    double lb = getLowerBound( _b );
+
+    // cast _b and _f to doubles
+    double reluInput = _tableau->getValue( _b );  // ReLU input before activation
+    double reluOutput = _tableau->getValue( _f ); // ReLU output after activation
+
+    // compute ReLU score
+    double scaler = ub / ( ub - lb );
+    double term1 =
+        std::min( scaler * reluInput * biasTerm, ( scaler - 1.0 ) * reluInput * biasTerm );
+    double term2 = ( scaler * lb ) * reluOutput;
+
+    return term1 - term2;
+}
+
+double ReluConstraint::calculateBias() const
+{
+    auto it = _biasCache.find( this );
+    if ( it != _biasCache.end() )
+        return it->second;
+
+    throw NLRError( NLRError::RELU_NOT_FOUND, "Bias not found in cache." );
+}
+
 double ReluConstraint::computePolarity() const
 {
     double currentLb = getLowerBound( _b );
@@ -1052,6 +1108,11 @@ PhaseStatus ReluConstraint::getDirection() const
     return _direction;
 }
 
+void ReluConstraint::updateScoreBasedOnBaBsr()
+{
+    _score = std::abs( computeBaBsr() );
+}
+
 void ReluConstraint::updateScoreBasedOnPolarity()
 {
     _score = std::abs( computePolarity() );

src/engine/ReluConstraint.h

@@ -37,6 +37,15 @@
 
 #include <cmath>
 
+/*
+  Namespace for Network Level Reasoner Instance.
+*/
+namespace NLR {
+class NetworkLevelReasoner;
+class Layer;
+struct NeuronIndex;
+} // namespace NLR
+
 class ReluConstraint : public PiecewiseLinearConstraint
 {
 public:
@@ -57,6 +66,21 @@ class ReluConstraint : public PiecewiseLinearConstraint
     */
     PiecewiseLinearConstraint *duplicateConstraint() const override;
 
+    /*
+      Setter for network level reasoner instance.
+    */
+
+    void setNetworkLevelReasoner( NLR::NetworkLevelReasoner *nlr )
+    {
+        _networkLevelReasoner = nlr;
+        initializeBiasCache( *_networkLevelReasoner );
+    }
+
+    /*
+      Cache biases for the source layers for ReLU neurons.
+    */
+    static void initializeBiasCache( NLR::NetworkLevelReasoner &nlr );
+
     /*
       Restore the state of this constraint from the given one.
     */
@@ -215,6 +239,10 @@ class ReluConstraint : public PiecewiseLinearConstraint
 
     bool supportPolarity() const override;
 
+    bool supportBaBsr() const override;
+
+    double computeBaBsr() const;
+
     /*
       Return the polarity of this ReLU, which computes how symmetric
       the bound of the input to this ReLU is with respect to 0.
@@ -227,6 +255,7 @@ class ReluConstraint : public PiecewiseLinearConstraint
       always between -1 and 1. The closer it is to 0, the more symmetric the
       bound is.
     */
+
     double computePolarity() const;
 
     /*
@@ -237,12 +266,18 @@ class ReluConstraint : public PiecewiseLinearConstraint
 
     PhaseStatus getDirection() const;
 
+    /*
+      Update the score based on the BaBsr heuristic
+    */
+    void updateScoreBasedOnBaBsr() override;
+
     void updateScoreBasedOnPolarity() override;
 
     const List<unsigned> getNativeAuxVars() const override;
 
 private:
     unsigned _b, _f;
+    NLR::NetworkLevelReasoner *_networkLevelReasoner;
     bool _auxVarInUse;
     unsigned _aux;
 
@@ -276,6 +311,12 @@ class ReluConstraint : public PiecewiseLinearConstraint
      Assign a variable as an aux variable by the tableau, related to some existing aux variable.
     */
     void addTableauAuxVar( unsigned tableauAuxVar, unsigned constraintAuxVar ) override;
+
+    /*
+      cached biases for the source layers for ReLU neurons.
+    */
+    static std::unordered_map<const ReluConstraint *, double> _biasCache;
+    double calculateBias() const;
 };
 
 #endif // __ReluConstraint_h__

src/engine/SnCDivideStrategy.h

@@ -23,6 +23,7 @@ enum class SnCDivideStrategy {
     // Relu splitting
     Polarity, // Pick the ReLU with the polarity closest to 0 among the first K nodes
     EarliestReLU,
+    BaBsrHeuristic, // Pick the ReLU with the highest BaBSR score
 
     Auto
 };