Reinforcement Learning for Neural Model Editing

Editing pretrained neural networks requires specialized algorithms tailored to specific objectives. Designing such algorithms is often time-consuming and demands significant effort. Motivated by the works of Li et al., 2016 and Ha et al., 2016, we investigate: Can reinforcement learning agents learn policies for modifying neural networks using only reward feedback? We present a framework that formulates neural model editing as a reinforcement learning problem, where agents modify models using reward feedback. We introduce two environments for an agent to interact with the parameters of neural networks. The reward function combines a utility-preservation objective with a task-specific editing objective. We evaluate the framework on bias mitigation in text classification and machine unlearning in image classification, both of which traditionally rely on specialized algorithms. **Environment: MaskWorld.** The agent has the ability to scale the weights by a constant. The updated weights for layer $l$ at timestep $t$ are given by ${W}^{(l)}_t = {W}^{(l)} \odot a_t$, where $a_t$ is the action at timestep $t$ and has the same dimensions as ${W}^{(l)}$. **Environment: ShiftWorld.** The agent has the ability to shift the weights by a constant. The updated weights for layer $l$ at timestep $t$ are given by ${W}^{(l)}_t = {W}^{(l)} + a_t$, where $a_t$ is the action at timestep $t$ and has the same dimensions as ${W}^{(l)}$. **Reward Function.** We define the reward function as a weighted combination of two components: (i) an utility term that captures the general performance of the model, and (ii) an editing objective-specific term that quantifies the extent to which the modified model satisfies a desired property. The reward function is defined as $R = \alpha U(w) + \beta O(w)$, where $U(w)$ denotes the utility of the model, $O(w)$ denotes the quantity associated with the target objective, and $\alpha$ and $\beta$ coefficients control the trade-off between the two components. **Machine Unlearning.** We train a CNN on the MNIST dataset, which serves as the original model. We task the agent with forgetting the class 7 while preserving performance on classes 1–6 and 8–9. **Bias Mitigation.** We train a neural network on the Jigsaw Toxic Comment Classification Challenge dataset (cjadams et al., 2017) to create a deliberately biased model. We task the agent with reducing biases acquired during training while preserving predictive performance on the original classification task. **Results.** Table 1 presents the performance of the learned policies on the machine unlearning task. We measured the accuracy on the retain set (classes 1-6, 8-9) and the forget set (class 7). The results demonstrate that policies learned are effective at removing information related to the forget set while preserving performance on the retain set. | Model/Method | Accuracy on Retain Set | Accuracy on Forget Set | |:---|:---:|:---:| | Original Model | $0.9241 \pm 0.0168$ | $0.9073 \pm 0.0433$ | | MaskWorld | $0.9375 \pm 0.0066$ | $0.0000 \pm 0.0000$ | | ShiftWorld | $0.9381 \pm 0.0060$ | $0.0001 \pm 0.0003$ | | Retrained from Scratch (Retain Set) | $0.9339 \pm 0.0071$ | $0.0000 \pm 0.0000$ | _Table 1: Higher retain and lower forget accuracy indicate better unlearning._ Table 2 reports the accuracy on test sets derived from the Jigsaw Toxic Comment Classification Challenge dataset (cjadams et al., 2017) (general dataset) and the synthetic test set from (Dixon et al., 2018) (unbiased dataset). Accuracy on the Jigsaw test set reflects the model’s general utility, while accuracy on the unbiased dataset serves as a measure of bias removal; higher accuracy on the unbiased dataset indicates a reduction in model bias. The results show that the modified model improves accuracy on the unbiased dataset by 7\%, indicating that the policy identified and adjusted weights associated with spurious correlations. | Model/Method | Accuracy on General Dataset | Accuracy on Unbiased Dataset | |:---|:---:|:---:| | Original Model | $0.7967 \pm 0.0091$ | $0.5447 \pm 0.0177$ | | MaskWorld | $0.7472 \pm 0.0365$ | $0.6171 \pm 0.0348$ | | ShiftWorld | $0.7490 \pm 0.0232$ | $0.6220 \pm 0.0308$ | | Fine-tuning on Unbiased Dataset | $0.7369 \pm 0.0313$ | $0.6061 \pm 0.0253$ | _Table 2: Higher accuracy on the unbiased dataset with minimal degradation on the general dataset indicates more effective bias mitigation._