Learning Aligned Stability in Neural ODEs Reconciling Accuracy with Robustness

This paper introduces Zubov-Net, a novel framework for Neural ODEs that reconciles accuracy and robustness by unifying dynamics and classification through learnable Lyapunov functions and a Zubov-driven mechanism to actively align prescribed regions of attraction with true decision boundaries.

The Gist

The Big Problem: The "Rigid Box" vs. The "Messy World"

Imagine you are teaching a robot to sort mail. You want it to be accurate (put every letter in the right bin) and robust (keep sorting correctly even if someone shakes the table, throws dust on the letters, or tries to trick it with a fake stamp).

Current AI models (like Neural ODEs) are great at sorting mail, but they are fragile. If you nudge a letter slightly, the robot might panic and drop it in the wrong bin.

To fix this, previous researchers tried to force the robot to be "stable." They built rigid, pre-made boxes (called Regions of Attraction) around each bin. They told the robot: "No matter what happens, if a letter is inside this box, it must stay in this box."

The Flaw: The real world is messy. The "boxes" the researchers built didn't match the actual shape of the mail piles.

• Sometimes the box was too small (the robot couldn't sort a slightly tilted letter).
• Sometimes the boxes overlapped (a letter was in two boxes at once, confusing the robot).
• Result: To make the robot safe, they had to make the boxes so rigid that the robot became slow and inaccurate. They traded accuracy for safety.

The Solution: Zubov-Net (The "Smart, Shape-Shifting Guard")

The authors propose a new framework called Zubov-Net. Instead of building rigid, pre-made boxes, they teach the robot to learn the shape of the boxes itself while it learns to sort the mail.

Here is how they do it, broken down into three simple steps:

1. The "Shape-Shifting" Classifier (The Unified Architect)

In old models, the "sorting logic" (the classifier) and the "safety rules" (stability) were two different people talking to each other. They often disagreed.

In Zubov-Net, the safety rule IS the sorting logic.

• Analogy: Imagine a security guard who doesn't just check IDs; the guard is the ID card. If the guard says "You belong in the VIP room," it's because the guard's internal sense of safety and the VIP room's location are perfectly synced.
• How it works: They use a special type of math function (a Lyapunov function) that acts as both the decision-maker and the safety net. This ensures the "safe zone" always matches the "correct answer."

2. The "Rubber Band" Alignment (Zubov-Driven Matching)

Even if the safety rule is the decision-maker, the robot might still drift. The authors use a famous math theorem (Zubov's Theorem) to create a rubber band between the "intended safe zone" and the "actual path the data takes."

• Analogy: Imagine you are walking a dog on a leash.
  • Old way: You force the dog to walk in a perfect square, even if the dog wants to go in a circle. The dog gets frustrated (inaccurate).
  • Zubov-Net way: You hold a rubber band. You tell the dog, "Stay in the circle I'm drawing." If the dog wanders, the rubber band gently pulls it back, but the circle itself can stretch and shrink to fit the dog's natural path.
• The Magic: They turn this math theorem into a "loss function" (a scorecard). If the robot's path drifts away from the safe zone, the score gets bad, and the robot learns to pull the safe zone to match the path, or pull the path to match the zone. They align perfectly.

3. The "Active Sculptor" (Controlling the Geometry)

Once the zones are aligned, the authors don't just leave them alone. They actively sculpt the zones to be as far apart from each other as possible.

• Analogy: Imagine you are arranging furniture in a room. You don't just put the sofa and the TV in the same corner. You actively push the sofa away from the TV to create a wide, clear path between them.
• The Result: By pushing the "safe zones" for different classes (e.g., "Cat" vs. "Dog") far apart, it becomes very hard for a small nudge (noise) or a trick (adversarial attack) to push a letter from the "Cat" zone into the "Dog" zone.

The Secret Sauce: The "PIACNN" (The Focused Lens)

To make all this work, they built a special neural network called PIACNN.

• The Problem: Making a network that is both "convex" (mathematically safe) and "smart" (can tell complex things apart) is hard. Usually, you have to choose one or the other.
• The Fix: They added an Attention Mechanism. Think of this as a spotlight. The network shines a light only on the features that matter for stability (like the shape of a cat's ear) and ignores the noise (like the background color). This keeps the math stable while making the AI very sharp.

Why This Matters (The Results)

The paper tested this on famous image datasets (like CIFAR-10 and Tiny-ImageNet).

• Clean Accuracy: It sorted the mail perfectly when the room was quiet.
• Robustness: When they shook the table (added noise) or tried to trick the robot (adversarial attacks), Zubov-Net kept sorting correctly, while other models started making mistakes.
• Efficiency: It didn't slow down the robot much. It was almost as fast as the standard models.

Summary in One Sentence

Zubov-Net stops forcing AI into rigid, mismatched safety boxes and instead teaches the AI to grow its own safety zones that perfectly match its decisions, making it both super smart and super tough against tricks.

Technical Summary

1. Problem Statement

Neural Ordinary Differential Equations (Neural ODEs) exhibit inherent robustness compared to discrete-depth networks, yet existing methods for enforcing stability often fail to reconcile accuracy with robustness.

• The Core Limitation: Current approaches impose rigid stability conditions (e.g., fixed Regions of Attraction (RoAs), global projection mechanisms, or per-sample equilibrium constraints). These methods create a misalignment between the model's prescribed stability regions (where the system is mathematically stable) and its logical decision boundaries (where classification occurs).
• The Consequence: This misalignment forces a trade-off: enforcing strict stability often shrinks inter-class margins, reducing discriminative power and accuracy. Conversely, optimizing purely for accuracy leaves the model vulnerable to adversarial attacks and noise.
• The Gap: Existing methods lack a mechanism to dynamically adapt stability regions to the learned data distribution, leading to over- or under-constrained systems that cannot effectively handle complex, multi-class discriminative structures.

2. Methodology: Zubov-Net

The authors propose Zubov-Net, a novel framework that unifies dynamics and decision-making to actively control the geometry of stability regions.

A. Unified Architecture: The Lyapunov Classifier

Instead of separating the feature dynamics from the classifier, Zubov-Net employs learnable Lyapunov functions directly as the multi-class classifier.

• Dual Role: For each class $c_i$, a Lyapunov function $W_i$ defines the Prescribed Region of Attraction (PRoA). The classifier predicts the class with the lowest Lyapunov value (closest to equilibrium).
• PIACNN (Partially Input-Attention-based Convex Neural Network): To implement this, the authors design a specialized network architecture.
  • It combines the convexity required for Lyapunov stability with an input-attention mechanism to enhance discriminative power.
  • The attention mechanism focuses on equilibrium-relevant features and acts as a weight normalization to prevent output explosion in deep architectures.
  • The network is proven to be convex with respect to the input, ensuring valid Lyapunov properties.

B. Zubov-Driven Region Matching

To bridge the gap between the prescribed stability (PRoA) and the true dynamical stability (RoA) of the Neural ODE, the authors reformulate Zubov's Theorem.

• Consistency Loss ($\mathcal{L}_{con}$): They transform Zubov's partial differential equation (PDE) into a differentiable consistency loss. This loss penalizes the discrepancy between the time derivative of the Lyapunov function along the ODE trajectory and the theoretical decay rate required for stability.
• Mechanism: Minimizing this loss forces the true dynamical RoAs to align with the PRoAs defined by the classifier, ensuring that the system's physical stability matches its logical decision regions.

C. Active RoA Control

Once alignment is established, the framework actively shapes the geometry of the attraction basins to maximize robustness.

• Tripartite Loss Function: The training objective consists of three components:
  1. Consistency Loss ($\mathcal{L}_{con}$): Ensures PRoA-RoA alignment (Zubov-driven matching).
  2. Classification Loss ($\mathcal{L}_{cla}$): Ensures trajectories converge to the correct class equilibrium (certifying asymptotic stability).
  3. Separation Loss ($\mathcal{L}_{sep}$): Actively maximizes the distance between PRoAs of different classes to prevent overlap and enlarge inter-class margins.
• Parallel Boundary Sampling: To compute the separation loss efficiently, the authors introduce an algorithm that samples boundary points between classes in parallel, avoiding expensive global optimization.

3. Key Contributions

1. Zubov-Net Framework: A unified architecture where learnable Lyapunov functions serve as multi-class classifiers, eliminating the structural misalignment between stability and decision boundaries.
2. Zubov-Driven Matching Mechanism: A novel method to align prescribed and true regions of attraction by reformulating Zubov's equation into a differentiable consistency loss.
3. Active RoA Control Paradigm: A shift from passive stability certification to active geometric control, optimizing the shape of attraction basins to reconcile accuracy and robustness.
4. PIACNN Architecture: A new network design that balances convexity (for stability) and expressiveness (for classification) via an input-attention mechanism.
5. Theoretical Guarantees:
   • Proof that minimizing the tripartite loss guarantees PRoA-RoA alignment, non-overlapping basins, and trajectory containment.
   • Derivation of a certified robustness margin ($\epsilon^*$) based on Lipschitz constants, linking basin geometry directly to robustness against norm-bounded perturbations.
   • Stochastic convex separability analysis justifying the use of convex classifiers in high-dimensional spaces.

4. Experimental Results

The framework was evaluated on SVHN, CIFAR-10, CIFAR-100, and Tiny-ImageNet.

• Accuracy & Robustness Trade-off: Zubov-Net achieved the highest clean accuracy and the highest average robustness against eight types of stochastic noise (e.g., Gaussian, motion, brightness) across all datasets.
• Adversarial Robustness:
  • Outperformed state-of-the-art stability-based baselines (LyaNet, Proj-NODE, SODEF) against white-box attacks (FGSM, PGD, BIM, APGD, Jitter).
  • Demonstrated superior performance against black-box attacks (VNI, Square, AutoAttack).
  • Showed strong synergy when combined with adversarial training (TRADES), achieving the best overall robustness scores.
• Geometric Evidence: t-SNE visualizations revealed that Zubov-Net maintains compact intra-class clusters and distinct inter-class boundaries even under perturbation, whereas baselines showed significant feature scattering.
• Efficiency: While training is slightly slower than standard Neural ODEs due to boundary sampling, Zubov-Net is significantly faster at inference than projection-based methods (Proj-NODE) and LyaNet, as it requires no stabilization operations during the forward pass.

5. Significance

This paper addresses a fundamental bottleneck in robust deep learning: the conflict between stability guarantees and classification performance.

• Paradigm Shift: It moves beyond "certifying" stability as a post-hoc constraint to "learning" stability as an integral part of the decision-making process.
• Theoretical-Practical Bridge: By leveraging Zubov's theorem (a classical control theory concept) in a differentiable, data-driven manner, the authors provide a rigorous theoretical foundation for robust Neural ODEs that is practically implementable.
• Generalizability: The proposed "aligned stability" framework offers a new direction for designing robust architectures not just for image classification, but potentially for time-series forecasting, reinforcement learning, and other dynamical systems where safety and reliability are critical.