Thank you for the feedback.

Key idea of co-modulation: we have revised the methods text to make clear what the starting point of the is from a biological perspective and given the stochastic modulation recent theory, and what are novel elements unique to our solution. Briefly, the previous work had proposed a mechanism for changing the readout of a sensory representation in a task specific way, based on task-relevance tags provided by targeted stochastic modulation. What is completely new in our paper is 1) the controller that dynamically decides which neurons to stochastically comodulate as a function of the task and 2) the nature of the computation: everything in the work before involved one task at the time whereas here we focus on different versions of multitask learning, in which the same inputs need to be classified along a large set of decision axes (in CelebA) or representations built for coarse categorization needs to be refines for high precision discriminations (in Cifar-100). Finally the computational relevance of comodulation has only ever been demonstrated in very simple toy tasks, so our work provides an opportunity to test the computational limits of the idea in more substantial machine learning problems (see also reply to reviewer 8Tjw).

Typos and co: thanks for catching those. We use ‘comodulation’ as it is shorter, but ‘co-modulation’ would also be fine. Color scheme: we have now revised all the figures to use a published color-blind visible palette. Hopefully that helps improve visual clarity.

Itemized answer to questions:

1. Similarities: the functional form of the gain modulation in the encoder layer is identical, the decoder gain computation procedure is in the same spirit, although different in the details. Differences: architecture, in particular the presence of the controller, nature of the tasks, the stages of the training procedure (see also general reply above). The insights about the computational roles of comodulation for task flexibility and the mechanistic insights about how it is achieved at the decision stage are also unique to our work.

2. Using ImageNet and COCO: The main purpose of comodulation is to tweak representations in a network trained on multiple tasks, so as to make the representations task-dependent. We thus need datasets that have an intrinsic multi-task logic to them. This is not the case for either ImageNet or COCO, at least not in their regular use. ImageNet would be thought of as a scaled up version of Cifar100, but without the coarse to fine labeling. COCO does have multiple tasks, but they are usually not learned together, and when they are, the models used have a large amount of unshared parameters. Our decoder formulation would likely not be useful in that setup.

3. Ours is not strictly a continuous learning setup, but it does have a potential advantage in that it structurally segregates common useful regularities in the data (in the backbone) from unique task specific adaptations (in the controller). It is not inconceivable that this separation could provide more robustness to across task interference, but we have not explicitly quantified this effects. It is however an idea worth further thinking on.

4. LoRA is low-rank adaption of large language models (https://arxiv.org/abs/2106.09685) , control net (https://arxiv.org/pdf/2302.05543.pdf) is a method for adapting diffusion models (hopefully we got the references you intended here). They mark bigger departures from the type of task adaptation intended here. Similarities: Both aim to fine-tune an existing representation. Differences: Lora optimizes new sets of weights which linearly modify the frozen weights in the attention heads, while control net learns new convolutional layers which modify the internal representations in some layer of the UNet. Both methods modify the representations at multiple stages of processing, whereas ours only directly modifies one layer. Neither are meant for multi-task learning, which is the main focus in this paper.

We thank the reviewer for the feedback.

Text clarity: We agree that rushing towards the deadline made the text fall short in terms of clearly explaining some of the details. We have now revised the main text to hopefully address various clarity concerns.

Training procedure: We tested both training scenarios: 1) where the controller is trained jointly with the network, and 2) where the controller is trained separately with the fixed pretrained backbone. We chose to report the second option for consistency with the other experiment (CIFAR-100 setup) and because we found that training the controller separately made the training more stable. Finally, training the controller is fast (you only need to backpropagate gradients up to the encoder) so that overall training time is not substantially increased if done in two stages.

Evaluation metrics: The metrics used were chosen based on the previous literature on multitask learning in CelebA, which are also reported in Table 1. Precision is defined as true positives/(false positives +true positives), thus if the network classifies only one instance as a positive and is correct it will have a precision of 1. Comodulation behaves a little unusual by this metric, for reasons that we don’t completely understand. Recall measures the amount of true positives/(false negatives + true positives); the two naturally trade-off between each other. The F1 score is the harmonic mean of the two, and quantifies the trade off between recall and precision. F1 is the key metric to look at. We don’t see any reason why the precision properties of comodulation would trivially explain the benefits seen in terms of F1 scores.

Itemized question answers: a) Separate phases of training are not typical, it is a design choice that we came to via numerical experimentation (see also training explanation above).

b) Confidence calibration is defined as model predicted confidence matching actual errors made for that with that confidence level. So for the set of images that the model predicts that it should be 0.4 sure about the class label, 40% of those should be mislabeled. We tried to rephrase the text to better convey the idea.

c) Code is only available for some options (Max roaming, Pascal et al. (2021)) so we used the evaluations previously reported in (ETR-NLP Ding et al. (2023)) Table 4, while matching our evaluation procedure to theirs as much as possible. We use the same batch size, augmentations and optimizer. Worth noting that the code is not available for ETR-NLP Ding et al. (2023), so one cannot be sure that the comparison is completely one-to-one. Nonetheless, we did our best to ensure it would be the case based on the methods described in the paper.

d) Thanks for pointing this out. The attention baseline is indeed the model that you describe, now explained in the text. The procedure for hyperparameter tuning was identical for both stochastic and deterministic attention. So, yes it did receive the same degree of attention (pun intended). The hyperparameter set included the size of the hidden layer of the controller in the range [32,64,128,256] and the learning rate of the controller, on a log scale from [0.0001 to 0.1]. The pretraining hyperparameters optimization aimed to achieve the best results on the pretraining task in the absence of the controller, as would be the case in practical ML applications.

e) Apologies, we referenced the wrong equation. What we meant is that the encoder is not stochastically modulated at test time, i.e we only use the stochastic modulator to estimate the decoding gain. We have changed the text to make this clearer.

Typos fixed, thank you for catching those.

We thank the reviewer for the feedback.

Regarding the magnitude of performance improvements due to comodulation: While it is true that the improvements diminish with increasing training epochs, the results from comodulation and attention are quite different when placed in the context of the previous literature. The attention model F1 sits at 67.8 which is below the previous SOTA of Ding et Al 2023. On the other hand, the comodulation has a F1 of 69.6, which in terms of increases relative to the previous SOTA is a significant jump. To put things in perspective, the two most recent increases in SOTA on CelebA were 0.5 and 0.8, which compared to our 1.8 are much smaller.

We apologize for the lack of clarity of the text, we have revised it to better explain the methodology and the significance of the results. We have also corrected the typos and fixed the issues with improper formatting of references.

Attention/deterministic gain modulation uses the exact same architecture with a multiplicative gain effect on the encoder and no gain effect on the decoder. The procedure for training and hyperparameter optimizations are shared with stochastic modulation, to keep the comparison as fair as possible. We have clarified these points in the methods.

We thank the reviewer for the feedback.

With regards to the biological relevance of this work: While data from several experimental datasets show that task-dependent comodulation is present in brain activity (e.g., Rabinowits for V4, Haimerl for V1 and MT), the computational significance of these empirical observations has only been demonstrated indirectly, and using very simple tasks. Thus, the main neuroscience-relevant question we address in this paper is “What kind of computations can such a mechanism support? Is it only viable in single toy tasks or does it have enough computational power to help perception and decision making in ways that matter outside the confines of an experiment?” By proving that stochastic comodulation does actually provide computational gains in challenging machine learning tasks we are, at the very least, showing that it would be beneficial for the brain to use it. Our grounding in neural data is qualitative rather than quantitative, but arguably that is true for many successful models in computational neuroscience. Finally, although abstract the model does speak to some qualitative features of neural data. Even if the learning part of the process is not biologically plausible in the details, the architectural structure of the model can motivate the experimental investigation of the brain circuitry performing the function of the controller (putatively superior colliculus) and makes predictions about its modulatory effects on visual areas. One should also note that some of the mechanistic details that one would like to know about stochastic modulation are equally mysterious for the deterministic gain modulation supporting what we typically refer to as “attention” in the brain, so these limitations are by no means unique to our model. Hopefully, our modeling results will further spur investigation of the neural mechanisms supporting task flexibility.

In terms of links to the previous literature: we have revised the main text to hopefully spell out more clearly the similarities and differences between the solution presented here and other machine learning approaches (Related work section) and the previous work on biological roles of stochastic modulation (Introduction and Methods).

Visual clarity: we have changed the training range for Fig 1B to make it clearer that comodulation surpasses the pretraining baseline and converges faster than gain modulation. We have also changed the color scheme to make it more discriminable for the spectrum of the readers. Even when the effects look small on the scale of the range of metric values seen across learning, hopefully the summary statistics make the significance of those distinctions statistically clear.

With regards to unclarity of explanations and various typos: we have revised the text to address each of those concerns. We have also corrected the cartoon to accurately reflect the one-hot encoding of the tasks.

The updated methods clarify a little more the role of the modulator at both encoding and decoding time, and the temporal relevance of those processes. Briefly, the modulator injects multiplicative noise in the encoding layer which propagates together with the image information to the decoder layer. The correlations between neural activity and the same modulator change the layer’s gains, which ultimately affect the output. For simplicity, the modulator is iid gaussian although this can definitely be relaxed with no ill effect. We have also tweaked the table to hopefully make the results easier to read and interpret, and added more explanation about the metric in the main text.

Overall, even if the immediate biological relevance of the model is limited, we do present a biologically motivated algorithm that improves on state of the art on nontrivial machine learning tasks. Moreover, the idea is counterintuitive enough that one would not have tried something like it were it not for the biological observations. Such instances of synergistic interaction between neuroscience and machine learning are not that common, which is why we hope that our work is of general interest to the ICLR community.