# Policy Gradient-Driven Noise Mask

arXiv: https://arxiv.org/abs/2406.14568 · HTML: https://mlsanity.com/paper/2406.14568/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 22 references that are in the map.
- First posted on arXiv 29 Apr 2024.

## Builds on (most cited first, 13 of 13)

- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty](https://mlsanity.com/paper/1912.02781/) (2019)
- [AutoAugment: Learning Augmentation Policies from Data](https://mlsanity.com/paper/1805.09501/) (2018)
- [MedMNIST Classification Decathlon: A Lightweight AutoML Benchmark for Medical Image Analysis](https://mlsanity.com/paper/2010.14925/) (2020)
- [Towards Deep Neural Network Architectures Robust to Adversarial Examples](https://mlsanity.com/paper/1412.5068/) (2014)
- [Concept Generalization in Visual Representation Learning](https://mlsanity.com/paper/2012.05649/) (2020)
- [Reinforcement Learning Algorithms: An Overview and Classification](https://mlsanity.com/paper/2209.14940/) (2022)
- [Data Augmentation in Training CNNs: Injecting Noise to Images](https://mlsanity.com/paper/2307.06855/) (2023)
- [Adversarial Policy Gradient for Deep Learning Image Augmentation](https://mlsanity.com/paper/1909.04108/) (2019)
- [Navigating Noise: A Study of How Noise Influences Generalisation and Calibration of Neural Networks](https://mlsanity.com/paper/2306.17630/) (2023)
- [RL-LOGO: Deep Reinforcement Learning Localization for Logo Recognition](https://mlsanity.com/paper/2312.16792/) (2023)

## Built on by (most cited first, 2 of 2)

- [Cross-D Conv: Cross-Dimensional Transferable Knowledge Base via Fourier Shifting Operation](https://mlsanity.com/paper/2411.02441/) (2024)
- [Distributional sentiment modeling and anomaly detection for consumer complaint assessment](https://mlsanity.com/paper/2609.31653/) (2026)

## Abstract (arXiv)

Deep learning classifiers face significant challenges when dealing with heterogeneous multi-modal and multi-organ biomedical datasets. The low-level feature distinguishability limited to imaging-modality hinders the classifiers' ability to learn high-level semantic relationships, resulting in sub-optimal performance. To address this issue, image augmentation strategies are employed as regularization techniques. While additive noise input during network training is a well-established augmentation as regularization method, modern pipelines often favor more robust techniques such as dropout and weight decay. This preference stems from the observation that combining these established techniques with noise input can adversely affect model performance. In this study, we propose a novel pretraining pipeline that learns to generate conditional noise mask specifically tailored to improve performance on multi-modal and multi-organ datasets. As a reinforcement learning algorithm, our approach employs a dual-component system comprising a very light-weight policy network that learns to sample conditional noise using a differentiable beta distribution as well as a classifier network. The policy network is trained using the reinforce algorithm to generate image-specific noise masks that regularize the classifier during pretraining. A key aspect is that the policy network's role is limited to obtaining an intermediate (or heated) model before fine-tuning. During inference, the policy network is omitted, allowing direct comparison between the baseline and noise-regularized models. We conducted experiments and related analyses on RadImageNet datasets. Results demonstrate that fine-tuning the intermediate models consistently outperforms conventional training algorithms on both classification and generalization to unseen concept tasks. https://github.com/convergedmachine/Policy-Gradient-Driven-Noise-Mask