Joint Explainability-Performance Optimization With Surrogate Models for AI-Driven Edge Services
Foivos Charalampakos, Thomas Tsouparopoulos, Iordanis Koutsopoulos
TL;DR
This work tackles the challenge of balancing predictive accuracy with explainability for AI at the network edge by jointly training a black-box model and an interpretable surrogate within a bi-level multi-objective optimization framework. It employs gradient-based MGDA to compute a descent direction that simultaneously reduces the black-box prediction error $\mathcal{L}_{\text{pred}}$ and the surrogate fidelity loss $\mathcal{L}_{\text{PF}}$, optimizing over $\boldsymbol{\theta}$ and $\boldsymbol{\phi}$ under a nested constraint. Empirically, the proposed approach yields substantial Fidelity gains (over $99\%$ improvement in Global Fidelity and up to $98.5\%$ in Neighborhood Fidelity) with only modest accuracy reductions, outperforming single-task and naive multi-task baselines across global and local explainability settings. This framework advances trustworthy edge AI by enabling more faithful surrogate explanations of complex models, with potential extensions to Federated Learning and broader explanation-quality metrics beyond fidelity.
Abstract
Explainable AI is a crucial component for edge services, as it ensures reliable decision making based on complex AI models. Surrogate models are a prominent approach of XAI where human-interpretable models, such as a linear regression model, are trained to approximate a complex (black-box) model's predictions. This paper delves into the balance between the predictive accuracy of complex AI models and their approximation by surrogate ones, advocating that both these models benefit from being learned simultaneously. We derive a joint (bi-level) training scheme for both models and we introduce a new algorithm based on multi-objective optimization (MOO) to simultaneously minimize both the complex model's prediction error and the error between its outputs and those of the surrogate. Our approach leads to improvements that exceed 99% in the approximation of the black-box model through the surrogate one, as measured by the metric of Fidelity, for a compromise of less than 3% absolute reduction in the black-box model's predictive accuracy, compared to single-task and multi-task learning baselines. By improving Fidelity, we can derive more trustworthy explanations of the complex model's outcomes from the surrogate, enabling reliable AI applications for intelligent services at the network edge.