Mask-Aware Aggregation in Neural Networks

• Mask-Aware Aggregation Rule is a method that uses explicit mask signals to guide the selection and weighting of features or parameter updates in neural networks.
• It enhances model privacy, robustness, and performance by suppressing non-aligned updates and reducing vulnerabilities like backdoor attacks.
• Applications include federated learning, computer vision, and secure aggregation, where masked updates optimize spatial focus and semantic alignment.

A mask-aware aggregation rule is a principled strategy for combining information in neural networks or distributed learning systems under explicit, spatial or semantic guidance by mask signals. Across recent literature, such rules appear in federated learning, vision-language modeling, semantic segmentation, video object detection, medical reporting, image restoration, and secure aggregation. Core to these methods is the modulation, selection, or protection of features or parameter updates according to mask structures, thereby enabling privacy, robustness, spatial focus, or semantic alignment.

1. Formal Definition and Variants

A mask-aware aggregation rule is any aggregation operator—typically in the form of a weighted sum, attention mechanism, or secure sum—where mask tensors guide selection, weighting, or masking of features, gradients, or model updates. In federated learning, masks may zero or scale parameter updates according to their class-relevance; in vision tasks, masks may spatially select regions of interest for pooling or attention; in privacy-preserving computation, per-element binary masks control which vector indices are included in the aggregate.

Examples span:

• Class-aware gradient masking in federated learning, wherein gradient components with large magnitude—measured on a server-side validation set for a client's dominant class—determine parameter-level masks for model update aggregation (Arazzi et al., 6 Mar 2025).
• Instance mask aggregation for temporal feature fusion in video object detection (Hashmi et al., 6 Dec 2024).
• Secure aggregation rules, where per-element masks ensure model updates are summed only where a minimum number of clients contribute, addressing data-reconstruction attacks (Suimon et al., 6 Aug 2025).
• Mask-guided cross-modal feature fusion in medical vision-LLMs (Maqbool et al., 31 Oct 2025), mural restoration (Lei et al., 10 Aug 2025), and segmentation (Shi et al., 2022, Ao et al., 2022, Jiao et al., 2023).

2. Class-Aware Masking in Federated Learning

The seminal “mask-aware aggregation rule” in federated learning operates via several key steps (Arazzi et al., 6 Mar 2025):

1. Class Assignment: Each client’s local model is evaluated over server validation data partitioned by class; the dominant class $c_i^*$ for client $i$ at round $t$ is $$\arg\max_c \mathrm{Accuracy}(\mathrm{LM}_i^t, V_{(c)})$$.
2. Gradient Masking: The per-parameter gradient $$G_i^t = \nabla_w \mathcal{L}(\mathrm{LM}_i^t; V_{c_i^*})$$ is computed. Parameters with $|G_i^t(\theta)|$ exceeding a client-specific $p$-th percentile threshold $\tau_i^t$ receive mask value 1; all others are scaled by $\gamma\ll1$, yielding $M_{i,\mathrm{raw}}^t$. Temporal smoothing produces the final mask $M_i^t$.
3. Mask Application: The masked update $$\tilde{\mathrm{LM}_i^t} = \mathrm{LM}_i^t \odot M_i^t$$ is sent.
4. Dynamic Weighting: Each update receives importance $\omega_i^t = \sum_{\theta} M_i^t(\theta)$; weights $w_i^t$ normalize these.
5. Model Aggregation: The global model is updated as

$$\mathrm{GM}^{t+1} = \sum_{i=1}^N w_i^t\,(\mathrm{LM}_i^t \odot M_i^t) = \left( \sum_{i=1}^N \omega_i^t (\mathrm{LM}_i^t \odot M_i^t) \right) / \left( \sum_{i=1}^N \omega_i^t \right)$$

1. Privacy Properties: Only gradient-derived masks and masked updates are visible to the server; dataset sizes and class distributions remain private. Robustness to non-IID data and backdoor attacks is obtained by suppressing gradients unaligned with class-specific learning objectives.

Compared to FedAvg, FedNova, and SCAFFOLD, this approach eliminates all reliance on client metadata, isolates class-relevant information, yields improved convergence on heterogeneous data (Dirichlet $\alpha=0.125,0.3,0.5$), and dramatically reduces backdoor attack success rates (>80% $\rightarrow$ <15%) (Arazzi et al., 6 Mar 2025).

3. Mask-Aware Aggregation in Computer Vision

In vision models, mask-aware aggregation rules modulate the flow or combination of features spatially, semantically, or temporally:

• Zero-Shot Segmentation: Mask-aware CLIP representations are obtained by inserting proposal-specific, mask-masked attention within the transformer blocks. Each mask proposal $M_n$ is embedded as a binary bias mask $B$ that restricts the class token's attention to the masked region, enabling each feature vector to encode only content local to that proposal. Mask-aware loss terms tie predicted segment class-scores to ground-truth IoUs with the mask, while self-distillation losses retain global zero-shot properties. The aggregation step is performed during masked attention:

$$F^{(i+1)*}_{cls} = \mathrm{Softmax}\left(\frac{QK^{T}}{\sqrt{d}} + B\right) V$$

where $B$ enforces spatial masking (Jiao et al., 2023).

• Video Object Detection: In the FAIM pipeline, mask-aware feature aggregation operates both at single-frame (via instance-masked convolutional features $M_{tj}$) and cross-temporal levels (via multi-head self-attention over stacked mask features):

$M_j^{agg} = \sum_{t=1}^m \alpha_{t,j}\,v_{t,j}$

with attention weights $\alpha_{t,j}$ computed from classification and mask feature projections (Hashmi et al., 6 Dec 2024). Mask-guided spatio-temporal aggregation reduces background noise and intra-class variance, yielding higher mAP at identical FPS compared to bounding-box-only pooling.

• Multimodal Medical Reporting: The PETAR-4B architecture fuses global and focal 3D PET/CT volume features, each patch-wise conditioned by addition with mask-embeddings. Outputs from full-volume and focal-crop streams are summed token-wise, pooled, and projected into the LLM space for report generation. Ablation reveals mask-aware streams significantly improve localization-focused report metrics (e.g., GREEN score, CIDEr) (Maqbool et al., 31 Oct 2025).

4. Secure and Privacy-Preserving Aggregation

In federated learning and privacy-sensitive regimes, mask-aware aggregation refers to protocol-level masking of model updates:

• Per-Element Secure Aggregation: Each client applies additional PRG-based masks only on vector entries where the local update is nonzero, keyed by client-decryptor secrets and PRF outputs. Upon aggregation, unmasking is permitted for a coordinate $k$ iff $|C[k]|\ge t$, ensuring the aggregate at $k$ does not reveal single-client contributions:

$$y[k] = \begin{cases} \sum_{i\in C} x_i[k], & |C[k]|\ge t \ \perp, & \text{otherwise} \end{cases}$$

Security is maintained against colluding servers, clients, or decryptors, and the mechanism remains modular atop existing SecAgg schemes (Suimon et al., 6 Aug 2025).

• AHSecAgg Protocol: Masking is performed by additive homomorphic expansion $$r_i = (r\cdot s_i, r^2\cdot s_i, ..., r^m\cdot s_i)$$; Shamir secret sharing enables dropout-tolerant unmasking. This is distinct for its efficiency (O(m+n) server cost, O(m+n²) client cost), absence of per-pair secret sharing, and security under both semi-honest and actively malicious models (Zhang et al., 2023).

5. Mask-Aware Aggregation in Image Restoration and Segmentation

Mask-aware feature fusion rules are prominent in structure-aware restoration and segmentation:

• Mural Restoration (CMAMRNet): Aggregation is implemented at two levels. (1) Mask-Aware Up/Down-Samplers (MAUDS) interleave mask channels with feature channels both during upsampling (channel alignment, depthwise fusion) and downsampling (interleaving, depthwise convolution), preserving mask sensitivity throughout resolution transitions. (2) Co-Feature Aggregator (CFA) multiplexes image and mask features via parallel focusing blocks, modulating texture with mask-derived attention and summing with residuals. Ablation demonstrates joint application improves PSNR, SSIM, MAE, and LPIPS compared to state-of-the-art mural inpainting (Lei et al., 10 Aug 2025).
• Few-shot Segmentation: MANet aggregates a fixed set of masks $M_i$ (e.g., $K=S^2$ for $S\times S$ grid), each with a predicted foreground probability $p_i$; the segmentation prediction is

$S(x, y) = \sum_{i=1}^K p_i M_i(x, y)$

This mask-classification approach yields state-of-the-art mIoU performance and requires no pixelwise direct correspondence (Ao et al., 2022). In the DCAMA model, the aggregation is performed via cross-attention between all query and support pixels, yielding

$\hat y_q = \sum_{s=1}^{N_s} \alpha_{q,s} M_s^s$

with attention weights determined by the similarity of deep features. Multi-scale and n-shot extensions are realized by stacking support pixels and applying the same aggregation in one pass, outperforming ensemble-based methods (Shi et al., 2022).

6. Theoretical and Practical Implications

Mask-aware aggregation mechanisms consistently prioritize spatial, semantic, or privacy-relevant structure during aggregation. In distributed learning, they enhance both privacy (by obviating the need for client-side metadata and by suppressing single-client leakage) and robustness (by prioritizing class-specific updates and suppressing adversarial or backdoor perturbations). In imaging tasks, mask-aware rules outperform heuristic region proposals, naively averaged features, or unstructured pooling, especially for fine-grained, multi-scale, or cross-modal fusion.

A plausible implication is that the mask-aware design paradigm unifies a class of approaches aiming to preserve localization, semantic consistency, or privacy during aggregation, under both collaborative (e.g., federated learning, multimodal reporting) and single-task (e.g., segmentation, restoration) settings. This suggests a wider applicability of such operators in both model- and protocol-level innovations.

7. Comparative Summary Table

Application Domain	Mask-Aware Rule Role	Representative Paper
Federated Learning	Gradient masking for privacy and robustness	(Arazzi et al., 6 Mar 2025)
Secure Aggregation	Per-element mask thresholding for privacy	(Suimon et al., 6 Aug 2025Zhang et al., 2023)
Vision-Language (PET/CT)	Embedding mask info for spatial grounding	(Maqbool et al., 31 Oct 2025)
Video/Object Detection	Instance mask-based temporal fusion	(Hashmi et al., 6 Dec 2024)
Image Restoration	Multi-stage mask-guided feature fusion	(Lei et al., 10 Aug 2025)
Few-Shot Segmentation	Mask-weighted aggregation of proposals/tokens	(Shi et al., 2022Ao et al., 2022Jiao et al., 2023)

• Privacy Preserving and Robust Aggregation for Cross-Silo Federated Learning in Non-IID Settings (Arazzi et al., 6 Mar 2025)
• Beyond Boxes: Mask-Guided Spatio-Temporal Feature Aggregation for Video Object Detection (Hashmi et al., 6 Dec 2024)
• Per-element Secure Aggregation against Data Reconstruction Attacks in Federated Learning (Suimon et al., 6 Aug 2025)
• AHSecAgg and TSKG: Lightweight Secure Aggregation for Federated Learning Without Compromise (Zhang et al., 2023)
• PETAR: Localized Findings Generation with Mask-Aware Vision-Language Modeling for PET Automated Reporting (Maqbool et al., 31 Oct 2025)
• CMAMRNet: A Contextual Mask-Aware Network Enhancing Mural Restoration Through Comprehensive Mask Guidance (Lei et al., 10 Aug 2025)
• Dense Cross-Query-and-Support Attention Weighted Mask Aggregation for Few-Shot Segmentation (Shi et al., 2022)
• Few-shot semantic segmentation via mask aggregation (Ao et al., 2022)
• Learning Mask-aware CLIP Representations for Zero-Shot Segmentation (Jiao et al., 2023)