Objective-Driven Differentiable Optimization of Traffic Prediction and Resource Allocation for Split AI Inference Edge Networks

Split AI inference partitions an artificial intelligence (AI) model into multiple parts, enabling the offloading of computation-intensive AI services. Resource allocation is critical for the performance of split AI inference. The challenge arises from the time-sensitivity of many services versus time-varying traffic arrivals and network conditions. The conventional prediction-based resource allocation frameworks have adopted separate traffic prediction and resource optimization modules, which may be inefficient due to discrepancies between the traffic prediction accuracy and resource optimization objective. This paper proposes a new, objective-driven, differentiable optimization framework that integrates traffic prediction and resource allocation for split AI inference. The resource optimization problem (aimed to maximize network revenue while adhering to service and network constraints) is designed to be embedded as the output layer following the traffic prediction module. As such, the traffic prediction module can be trained directly based on the network revenue instead of the prediction accuracy, significantly outperforming the conventional prediction-based separate design. Employing the Lagrange duality and Karush-Kuhn-Tucker (KKT) conditions, we achieve efficient forward pass (obtaining resource allocation decisions) and backpropagation (deriving the objective-driven gradients for joint model training) of the output layer. Extensive experiments on different traffic datasets validate the superiority of the proposed approach, achieving up to 38.85% higher network revenue than the conventional predictive baselines.