Paper: Task adaptation by biologically inspired stochastic comodulation (ICLR-2024)

Reviewer comment (Writing):
While it is interesting to see a bio-inspired algorithm, I found the motivation is a bit hard to follow without the background knowledge about "co-modulation" in the brain. There is somehow a gap between the neuroscientific findings and the proposed deep learning methodology in this paper (I have not read the preset papers (Haimerl et al. 2019; 2022)).

Author response:
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).

Reviewer comment (Novelty):
The novelty of this paper is unclear in presence of (Haimerl et al. 2022). See the questions below.

Author response:
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).

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.

Reviewer comment (Presentation):
Colors in Figure 1, 2, 3, 5 are not friendly to color-blind readers (especially red-green), consider using different marker/line styles instead.

Author response:
Color scheme: we have now revised all the figures to use a published color-blind visible palette. Hopefully that helps improve visual clarity.

Reviewer comment (Writing):
Is there any reason that Haimerl et al. 2022 used the term "co-modulation" but this paper use "comodulation"?

Author response:
Typos and co: thanks for catching those. We use ‘comodulation’ as it is shorter, but ‘co-modulation’ would also be fine.

Reviewer comment (Novelty):
The paper wrote "We extend the model of stochastic comodulation presented in Haimerl et al. (2019; 2022) ...". Could the authors explain in more details about the relationship between the current paper and (Haimerl et al. 2022), including methodology, task setting etc.?

Author response:
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.

Reviewer comment (Experiments):
I also wonder how the presented stochastic comodulation method perform on more challenging dataset such as ImageNet and COCO. Is there any reason that ImageNet was not tested?

Author response:
Itemized answer to questions:
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.

Reviewer comment (Novelty):
What are the relation between the proposed methods and LoRA / control net, as these methods freeze the pretrained-model and conducts adaptation with additional network of fewer parameters?

Author response:
Itemized answer to questions:
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.

Reviewer comment (Experiments):
**Experimental details seem ad-hoc** The training setup seems a little odd. As far as I can tell, for the CelebA experiment, the authors (1) took a pre-trained ImageNet backbone (2) trained the whole network on the full task (3) fine-tuned just the controller parameters. This makes it difficult to compare the “efficiency” of training because there are many stages involved. Can the controller not be trained at the same time as the encoder/decoder (and if so, is this due to training problems or due to some experimental design choices)?

Author response:
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.

Reviewer comment (Experiments):
Is it typical to first train the entire network on all tasks and then fine-tune it on the same dataset? I am most familiar with work that fine-tunes for new tasks, so just clarifying whether this is a standard choice would be helpful.

Author response:
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).

Reviewer comment (Writing):
I’m a bit confused by the sentence “In other words, when a network predicts class C with a probability of 0.4, the probability that the network is correct is 0.4.” in the “Comodulation improves confidence calibration” section on Page 7). Is there a typo? If not, could you explain this in more detail, as I don’t understand how this could be correct as written. For instance, if Class C is only ever encountered <0.01% of the time then the network will be correct with a probability significantly less than 0.4, right?

Author response:
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.

Reviewer comment (Experiments):
In Table 1, are the comparison “state of the art” methods computed on models that were trained in the same environment that is used in this paper, or are these taken directly from the cited papers? This detail seems important to clarify, as there might be other underlying differences.

Author response:
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.

Reviewer comment (Reproducibility):
Could the authors provide further explanation of the deterministic “attention” model? I could not find details of this in the paper. This seems particularly important to explain, given that it is one of the critical comparisons and does best on the “Precision” metric. Is this attention just the controller without a stochastic element, and if not, why is something along those lines not a direct comparison?

Author response:
d) Thanks for pointing this out. The attention baseline is indeed the model that you describe, now explained in the text.

Reviewer comment (Experiments):
As a followup to (d), did the tuning of hyperparameters for training etc. for this “attention” model receive as much tweaking as the co-modulation comparison? It is a little puzzling to me that there is very little change in the representations after fine-tuning here, and so I wonder if it is the stochastic co-modulation that is actually helping, or if the “attention” case had hyperparameters that did not behave well with the training setup.

Author response:
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.

Reviewer comment (Reproducibility):
On page 4 there is the sentence “...we only use stochasticity to compute the gain, but we use unperturbed representations for decision making.” This doesn’t make sense to me – isn’t the gain perturbing the representations? Or is stochasticity somehow turned off during the testing (and only on during training?)

Author response:
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.

Reviewer comment (Evaluation):
The improvements produced by comodulation over attention seem quite marginal and tend to diminish as the number of pretraining epochs increases. It is true that as backbone size increases, pretraining becomes less of a viable option, but the small gains over attention regardless (with few pretraining epochs) makes the work less exciting.

Author response:
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.

Reviewer comment (Reproducibility):
Although the attention baseline is repeatedly evaluated in many parts of the paper, I couldn't find a mathematical description of this approach, adding which would improve the clarity of this work.

Author response:
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.

Reviewer comment (Evaluation):
Final results might need more grounding to actual neural data to claim that this paper has made stochastic modulation seem more like it might be used in the brain. "Might be a computationally viable candidate" is technically accurate, but the language still feels overall overstated with respect to how much these results mean for the actual brain.

Author response:
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.

Reviewer comment (Presentation):
It really is hard to see in many of the figures why we should be excited about the differences in the curves. Annotating or zooming in to ground them quantitatively would help.

Author response:
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.

Reviewer comment (Presentation):
Figure 1: unclear how the stochastic modulator originates and interacts with the decoder. The controller has task info embedded, but Fig 1 + text makes the stochastic modulator look both like it comes from the controller and like it's purely randomly sampled. Which is it?

Author response:
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.

Reviewer comment (Presentation):
Table 1 would benefit from far more annotation - it's a lot to take in right now, and the text only refers to it as a whole (or it refers to specifics that are hard to find and hold in memory).

Author response:
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.

Reviewer comment (Presentation):
Minor: the paper says the task is given as a one-hot, but the figure makes it seem like a binary representation of some kind.

Author response:
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.

Reviewer comment (Experiments):
To some degree it's unclear which audience this paper is for - for neuroscientists, there's almost no quantitative grounding in neural data. That said, the paper seems primarily geared toward the ML community - here, more/larger-scale experimentation would help but the results do show clear and consistent improvement. I think this is a smaller but well-scoped paper that points to interesting further research along with making a contribution.

Author response:
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.

Reviewer comment (Reproducibility):
How exactly is the stochastic modulator signal integrated into the decoder, and how is it generated? Does it come from the controller or a totally random sample? How does that reflect on Fig 1?

Author response:
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.