Cost-aware adaptive defense against feature-space adversarial attacks in partitioned deep neural networks

Partitioned deep neural network inference exposes intermediate feature representations to potential manipulation during transmission between computational nodes. Adversarial perturbations injected into these boundary tensors can propagate through the remaining network and cause downstream misclassification, while existing defenses either detect anomalies without restoring prediction reliability or apply unconditional recovery at unnecessary computational cost. This paper presents a cost-aware adaptive defense framework for feature-space adversarial attacks in partitioned deep neural networks operating under realistic communication degradation in distributed edge AI systems. A lightweight detector models clean boundary representations using channel-wise feature statistics and identifies anomalous deviations with controlled false-positive behavior. For flagged samples, a repair module projects anomalous statistics toward the clean region and performs moment-matching correction of the received feature map, with optional masked repair to restrict modification to the most anomalous channels. An MD-gated contextual bandit supervisor then selects the recovery action to balance robustness and computational cost. Experiments on partitioned VGG16 and ResNet50 under channel-sparse Gaussian-noise perturbations, together with an additional iterative query-based feature-space attack based on Natural Evolution Strategies (NES), show that the proposed framework maintains near-clean accuracy, substantially recovers attacked performance, and reduces unnecessary recovery compared with unconditional repair. The results demonstrate that lightweight detection combined with selective and adaptive correction provides an effective defense for robust distributed inference.