OptiGradTrust: Robust Federated Learning

• OptiGradTrust is a Byzantine-robust federated learning framework that uses a six-dimensional gradient fingerprint and dual-attention reinforcement learning to evaluate client updates.
• It integrates adaptive trust assignment with FedBN-Prox, effectively mitigating malicious attacks and handling non-IID data across clients.
• Empirical results demonstrate notable improvements in accuracy and convergence on diverse datasets, making it particularly suited for sensitive applications like medical imaging.

OptiGradTrust is a Byzantine-robust federated learning framework that evaluates client gradient updates through a six-dimensional fingerprint, maps those signals to adaptive trust weights through a dual-attention and reinforcement-learning module, and combines this defense layer with FedBN-Prox to address severe data heterogeneity (Karami et al., 31 Jul 2025). In the cited literature, the explicit name refers to distributed training across private local datasets—especially in medical settings—under a threat model with arbitrary gradient manipulation, while nearby work uses related trust-centered ideas for optimization modeling, noisy trust-region methods, evolutionary search, and agentic tool governance under different names (Karami et al., 31 Jul 2025).

1. Problem setting and defining characteristics

OptiGradTrust is formulated in the standard federated learning setting with $N$ clients, each holding a private dataset $\mathcal{D}_k$, and a global objective

$$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$

where $F_k$ is the local loss (Karami et al., 31 Jul 2025). Training proceeds over $T$ communication rounds, with 5–8 local epochs per round, and only model updates or gradients are shared; raw data never leaves clients (Karami et al., 31 Jul 2025).

The framework targets two failure modes simultaneously. The first is Byzantine behavior: up to a fraction $f \le 0.3$ of clients per round can be malicious, potentially colluding and knowing the current global model (Karami et al., 31 Jul 2025). The second is statistical heterogeneity: non-IID client distributions can induce client drift and slow or destabilize convergence (Karami et al., 31 Jul 2025). OptiGradTrust addresses both by combining anomaly-aware gradient analysis, adaptive trust assignment, and a heterogeneity-aware optimizer (Karami et al., 31 Jul 2025).

The attacks explicitly modeled are scaling, partial scaling, sign-flipping, additive Gaussian noise, and label-flipping (Karami et al., 31 Jul 2025). Scaling multiplies gradients by a large factor such as $\lambda = 10$; partial scaling scales 50% of gradient coordinates by 5; sign-flipping multiplies gradients by $-1$; additive Gaussian noise uses $\sigma = 5$ on MNIST or $\sigma = 10$ on CIFAR-10 and MRI; and label-flipping flips labels with probability $\mathcal{D}_k$0 on malicious clients’ local training sets (Karami et al., 31 Jul 2025). Honest clients are assumed to be in the majority (Karami et al., 31 Jul 2025).

A defining feature of the framework is that trust is not treated as a fixed heuristic. Instead, the server constructs a per-client representation from the raw update and from six auxiliary signals, then assigns trust weights $\mathcal{D}_k$1 that directly modulate aggregation (Karami et al., 31 Jul 2025). This makes OptiGradTrust a trust-aware robust aggregation scheme rather than a static outlier filter.

2. System architecture and round-level workflow

At round $\mathcal{D}_k$2, the server broadcasts the current global model $\mathcal{D}_k$3 to all clients, and client $\mathcal{D}_k$4 initializes $\mathcal{D}_k$5 (Karami et al., 31 Jul 2025). Each client then performs FedBN-P local training, updating local Batch Normalization parameters but excluding them from cross-client aggregation, and optimizes the remaining parameters with Adam against a proximalized objective (Karami et al., 31 Jul 2025). The client computes

$\mathcal{D}_k$6

and sends $\mathcal{D}_k$7 to the server (Karami et al., 31 Jul 2025). Malicious clients may arbitrarily manipulate this update before transmission (Karami et al., 31 Jul 2025).

On the server side, each received update passes through a multi-stage pipeline. The Feature Extraction Module computes a six-dimensional fingerprint $\mathcal{D}_k$8 from $\mathcal{D}_k$9 (Karami et al., 31 Jul 2025). A Dual-Attention Network then processes both the raw high-dimensional gradient and the six derived features to produce an embedding $$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$0 (Karami et al., 31 Jul 2025). A Double DQN policy observes the set of client embeddings and historical context, chooses a discrete trust action $$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$1, and converts it into a trust weight $$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$2 (Karami et al., 31 Jul 2025). The global model is updated through trust-weighted aggregation: $$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$3 where $$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$4 is the set of participating clients in round $$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$5 (Karami et al., 31 Jul 2025).

The server pipeline also includes an optional Privacy Engine that applies gradient quantization through QSGD 4× and differential privacy noise, although the core OptiGradTrust results do not rely on differential privacy (Karami et al., 31 Jul 2025). In architectural terms, the framework is therefore a coupled defense-and-optimization system: feature extraction identifies suspicious or useful update properties, reinforcement learning converts those properties into dynamic trust, and FedBN-P stabilizes the underlying distributed optimization (Karami et al., 31 Jul 2025).

3. Six-dimensional gradient fingerprint

The core analytic object in OptiGradTrust is the gradient fingerprint $$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$6, designed to capture anomaly, similarity, magnitude, sign pattern, and game-theoretic utility in a single representation (Karami et al., 31 Jul 2025). The six signals are deliberately heterogeneous, with the stated rationale that they provide orthogonal signals that are hard to simultaneously spoof (Karami et al., 31 Jul 2025).

Feature	Definition or source	Intended signal
$$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$7	VAE reconstruction error	anomaly score of gradient shape
$$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$8	cosine similarity to trusted server reference gradient	directional alignment
$$\theta^* = \arg\min_{\theta} F(\theta) = \arg\min_{\theta} \sum_{k=1}^{N} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} F_k(\theta),$$9	average cosine similarity to other clients’ gradients	consensus alignment
$F_k$0	$F_k$1 norm of the gradient	magnitude anomaly
$F_k$2	sign-consistency ratio with the reference gradient	coordinate-wise agreement
$F_k$3	Monte Carlo Shapley value	marginal validation contribution

The first feature is VAE reconstruction error. A Variational Autoencoder is trained on historically observed benign gradients, with encoder layers $F_k$4, a latent space of dimension 64, and a mirrored decoder; it is updated every 20 rounds (Karami et al., 31 Jul 2025). Given a gradient $F_k$5, the VAE outputs $F_k$6, and the anomaly score is

$F_k$7

Large reconstruction error flags deviations from the learned manifold of benign updates (Karami et al., 31 Jul 2025).

The second and fifth features use a trusted server reference gradient $F_k$8, computed from a global validation set of 1000 samples (Karami et al., 31 Jul 2025). The cosine-alignment feature is

$F_k$9

while the sign-consistency feature is

$T$0

These are particularly informative against sign-flip and label-flip attacks (Karami et al., 31 Jul 2025).

The third feature measures agreement with peers: $T$1 Its naive complexity is $T$2 per round, but with 10 clients this is inexpensive in the reported experiments (Karami et al., 31 Jul 2025). The fourth feature is the gradient norm,

$T$3

which directly exposes scaling and partial-scaling attacks (Karami et al., 31 Jul 2025).

The sixth feature is a Monte Carlo approximation to the Shapley value of a client’s update with respect to validation performance (Karami et al., 31 Jul 2025). For coalition value $T$4, the exact Shapley value is

$T$5

and OptiGradTrust estimates it by averaging marginal contributions across $T$6 random permutations, with $T$7 by default, $T$8 in attack-dense rounds, and $T$9 in benign rounds (Karami et al., 31 Jul 2025). The resulting contribution estimate is exponentially smoothed: $f \le 0.3$0 This is the most explicitly utility-oriented feature in the fingerprint, because it evaluates how much a client’s update improves or degrades validation performance (Karami et al., 31 Jul 2025).

4. Dual attention and reinforcement-learning trust assignment

The trust assignment mechanism uses two attention streams. In the gradient stream, the high-dimensional update is split into $f \le 0.3$1 chunks, with $f \le 0.3$2 in the implementation description, and processed first by local multi-head attention and then by a global attention layer (Karami et al., 31 Jul 2025). The local attention takes the standard form

$f \le 0.3$3

and the global layer is

$f \le 0.3$4

with key dimension $f \le 0.3$5 (Karami et al., 31 Jul 2025). This stream is intended to capture fine-grained spatial patterns in the parameter update (Karami et al., 31 Jul 2025).

The feature stream operates over the six-dimensional fingerprint: $f \le 0.3$6 which yields attention weights over the six features (Karami et al., 31 Jul 2025). This allows the model to shift emphasis across attacks—for example, toward magnitude-related features under scaling or toward sign-consistency under sign-flipping (Karami et al., 31 Jul 2025). The two streams are fused by

$f \le 0.3$7

where $f \le 0.3$8 is ReLU (Karami et al., 31 Jul 2025).

Trust assignment is formulated as a sequential decision problem and implemented with Double DQN (Karami et al., 31 Jul 2025). The state $f \le 0.3$9 consists of the set of client embeddings together with history features such as moving averages of trust scores, prior detection statistics, and accuracy/loss trends (Karami et al., 31 Jul 2025). The action space is discrete: $\lambda = 10$0 interpreted directly as trust weights $\lambda = 10$1 (Karami et al., 31 Jul 2025). The reward is

$\lambda = 10$2

with $\lambda = 10$3, $\lambda = 10$4, $\lambda = 10$5, and $\lambda = 10$6 (Karami et al., 31 Jul 2025). Here $\lambda = 10$7 is validation accuracy improvement, FPR is false positive rate, FNR is false negative rate, and EFF is an efficiency or confidence measure favoring more decisive weights (Karami et al., 31 Jul 2025).

The Q-network is a fully connected network with layers $\lambda = 10$8, ReLU activations, and dropout 0.3 (Karami et al., 31 Jul 2025). Experience replay has size 1000, updates occur every 10 rounds with mini-batch size 64, the learning rate is $\lambda = 10$9, the discount factor is $-1$0, and $-1$1-greedy exploration decays from 0.3 to 0.05 over 100 rounds (Karami et al., 31 Jul 2025). The trust mechanism is therefore adaptive in two senses: it learns which features matter most under current attack patterns, and it learns how aggressively to trust or suppress each client in a way that trades off robustness and benign-client retention (Karami et al., 31 Jul 2025).

5. FedBN-Prox and optimization under heterogeneity

FedBN-P is the optimization layer beneath the trust mechanism. It combines Federated Batch Normalization and FedProx by keeping BatchNorm parameters local while applying a proximal regularizer to the shared non-BN parameters (Karami et al., 31 Jul 2025). The local client objective is

$-1$2

with $-1$3 and, in the reported implementation, $-1$4 chosen via validation-based grid search (Karami et al., 31 Jul 2025).

The motivation is specific to heterogeneous federated data. FedProx discourages local drift by penalizing deviation from the broadcast model, while FedBN accommodates per-client distribution shifts by preserving client-specific running statistics and affine BatchNorm parameters (Karami et al., 31 Jul 2025). In medical imaging, the intended interpretation is that per-hospital intensity and distribution differences are absorbed into local BN statistics, whereas higher-level convolutional structure remains shared (Karami et al., 31 Jul 2025).

Relative to standard baselines, FedBN-P differs from FedAvg by adding the proximal term and by not averaging BN parameters; it differs from FedProx by adopting the FedBN treatment of BN layers; and it differs from FedBN by adding proximal regularization (Karami et al., 31 Jul 2025). The authors do not present a new formal convergence theorem, but they report empirical convergence behavior on Alzheimer MRI in which FedBN-P reaches an average accuracy of 95.27% and converges in 26 rounds, compared with FedBN at 95.75% in 30 rounds, FedProx at 92.55% in 24 rounds, FedNova at 90.80% in 35 rounds, and FedAvg at 90.64% in 30 rounds (Karami et al., 31 Jul 2025). This places FedBN-P on the stated accuracy–convergence Pareto frontier (Karami et al., 31 Jul 2025).

6. Evaluation, baselines, and empirical results

The reported experiments use 10 clients, all participating in each round, over 25–30 global rounds with 5–8 local epochs per round (Karami et al., 31 Jul 2025). Optimization uses Adam with learning rate $-1$5 and weight decay $-1$6 (Karami et al., 31 Jul 2025). The models are a 3-layer CNN for MNIST, ResNet-18 for CIFAR-10, and ResNet-18 on $-1$7 images for the Alzheimer MRI dataset, which contains 6,983 synthetic scans across four classes (Karami et al., 31 Jul 2025). Data partitions include IID, Dirichlet non-IID with $-1$8 and $-1$9, label-skew with 70% or 90% dominant class, and quantity skew via a log-normal distribution of data sizes (Karami et al., 31 Jul 2025).

Under IID conditions, OptiGradTrust with FedBN-P achieves 99.41% clean accuracy on MNIST and 99.21% average accuracy across the five attacks; 83.90% clean accuracy on CIFAR-10 and 82.44% average under attack; and 97.24% clean accuracy on Alzheimer MRI and 96.61% average under attack (Karami et al., 31 Jul 2025). On Alzheimer MRI specifically, the attack-wise accuracies are 96.92% for scaling, 96.75% for partial scaling, 96.60% for sign-flip, 96.50% for noise, and 96.30% for label-flip (Karami et al., 31 Jul 2025). In IID settings, the accuracy drop due to attacks is described as very small, typically below 1–1.5 percentage points (Karami et al., 31 Jul 2025).

Under non-IID regimes, the average performance remains comparatively stable. For Dirichlet $\sigma = 5$0, the averages over attacks are 98.12% on MNIST, 81.36% on CIFAR-10, and 95.00% on Alzheimer (Karami et al., 31 Jul 2025). For Dirichlet $\sigma = 5$1, they are 97.42%, 79.42%, and 93.80%, respectively; under label skew 90%, they are 97.21%, 79.20%, and 93.26% (Karami et al., 31 Jul 2025). Relative to IID attack performance, these losses are reported as modest despite strong non-IID distortion plus attacks (Karami et al., 31 Jul 2025).

The principal baselines are FLGuard, FLTrust, and FLAME (Karami et al., 31 Jul 2025). In IID averages over five attacks, the reported results are 99.21 for OptiGradTrust versus 98.98 for FLGuard, 98.61 for FLTrust, and 97.08 for FLAME on MNIST; 82.44 versus 81.72, 81.05, and 79.60 on CIFAR-10; and 96.61 versus 95.89, 95.44, and 92.90 on Alzheimer (Karami et al., 31 Jul 2025). In Dirichlet $\sigma = 5$2 non-IID settings, the results are 98.12 versus 96.53, 95.85, and 94.63 on MNIST; 81.36 versus 79.62, 78.80, and 77.56 on CIFAR-10; and 95.00 versus 93.41, 92.81, and 91.11 on Alzheimer (Karami et al., 31 Jul 2025). The reported improvements over FLGuard under heterogeneity are therefore approximately $\sigma = 5$3 percentage points on MNIST, $\sigma = 5$4 on CIFAR-10, and $\sigma = 5$5 on Alzheimer MRI (Karami et al., 31 Jul 2025).

The computational bottleneck is Monte Carlo Shapley estimation, since it requires repeated incremental aggregation and validation evaluation across permutations (Karami et al., 31 Jul 2025). Even so, the full experiments were reported to run on a single RTX 3090 GPU and to complete each dataset in about seven hours (Karami et al., 31 Jul 2025). This performance profile is consistent with the framework’s intended deployment niche, where robustness and heterogeneity handling can justify substantial server-side computation (Karami et al., 31 Jul 2025).

7. Limitations, term ambiguity, and related research directions

The authors identify five main limitations: scalability with many clients, hyperparameter sensitivity, exposure to fully adaptive attackers that mimic benign multi-feature profiles, limited evaluation scale in terms of model size and client count, and potential mismatch with very tight real-time constraints (Karami et al., 31 Jul 2025). The scalability issue is tied especially to pairwise similarity and Monte Carlo Shapley components, which grow at least quadratically in the number of clients (Karami et al., 31 Jul 2025). Suggested future directions include scalable Shapley approximations, multi-agent RL or hierarchical trust modeling, extension to more complex domains such as 3D MRI and federated GANs, domain-specific priors in the attention modules, and stronger integration with secure aggregation and privacy mechanisms (Karami et al., 31 Jul 2025).

A frequent source of ambiguity is that the name “OptiGradTrust” can suggest other trust-centered optimization paradigms. In the supplied literature, however, the exact title denotes the federated learning framework described above (Karami et al., 31 Jul 2025). Nearby work uses related ideas under different names. “OptiTrust” is a trustworthy optimization modeling agent for LP and MILP problems built from verifiable synthetic data, teacher-trajectory filtering, multi-language inference, and majority-vote cross-validation rather than federated aggregation (Lima et al., 5 Aug 2025). A separate line of work studies a noise-robust trust-region method for smooth nonconvex optimization with bounded function and gradient errors by modifying the ratio test through

$\sigma = 5$6

which is structurally unrelated to Byzantine-robust FL even though the phrase “trust” appears in the algorithmic mechanism (Sun et al., 2022). In evolutionary computation, Trust-Based Optimization replaces fixed migration in island models with trust- or reputation-driven interaction among agent-islands (Urbańczyk et al., 29 Oct 2025). In agentic tool governance, trust calibration has been formalized as Gaussian-process preference learning over allow, block, and ask regions, with a gateway rule

$\sigma = 5$7

which addresses progressive autonomy rather than distributed optimization (Ou, 18 May 2026).

This suggests a broader conceptual neighborhood in which “trust” functions as a control variable over optimization or decision-making, but the specific encyclopedic referent of OptiGradTrust is the federated learning system that combines multi-feature gradient analysis, RL-based trust weighting, and FedBN-Prox to defend against Byzantine attacks under non-IID data (Karami et al., 31 Jul 2025).