Paper: Memory Efficient Neural Processes via Constant Memory Attention Block (ICLR-2024)

Reviewer comment (Evaluation):
The authors might be more explicit about the ways in which memory is a meaningful limitation in conventional Neural Process methods. It may be possible to more-explicitly show both memory and method performance in a plot or table to illustrate the interplay.

Author response:
Thank you for your suggestion! We have updated the introduction to make more explicit how memory is a limitation in Neural Processes. The memory gains and method performance can be inferred from Figure 3 and the table results. If it would help, we can prepare the plot for the camera-ready.

Reviewer comment (Experiments):
While the authors demonstrate the relationship between block size and end-to-end time cost (in the appendix), doing so for their proposed approach compared to other latent variable modeling approaches might be helpful to contextualize CMANPs in terms of not only memory, but also computation time.

Author response:
We have included plots in the Appendix (Figure 11) which provide several comparisons of the computation time (i.e., runtimes). 

In Figure 11a, we compare the runtime of the update process of CMAB-based NP (CMANPs) with Perceiver's iterative attention-based NP model (LBANPs) and a transformer-based NP model. We see that the CMAB-based model only requires a constant amount of time to perform the update. In contrast, Perceiver's iterative attention-based model's update runtime scales linearly and Transformer model's update runtime scales quadratically.

In Figure 11b, we compare the runtime of querying (performing inference) of CMANP with LBANP (Perceiver's iterative attention-based model), TNP (Transformer-based model). We see that CMANP and LBANP stay constant due to only needing a fixed number of latents to perform inference while the transformer-based model (TNP) scales quadratically in runtime.

Reviewer comment (Writing):
The authors may care to briefly explain here and in the manuscript: why is permutation invariance important for latent variable/neural process modeling?

Author response:
From a theoretical and historical perspective, Neural Processes (Garnelo et al., 2018a;b) were originally inspired by Gaussian Processes which are permutation invariant in the context. 

However, from a practical perspective, permutation invariance is a generally useful property for versatile models. For example, self-attention is by default permutation-invariant. As a result, Transformers have become widely applicable by augmenting the input with various modality-specific components such as a positional encoding for language modelling or patches of an image for computer vision. Similarly, NPs have been applied to temporal data and images, leveraging modality-specific components. We refer the reviewer to the survey by Jha et al., 2023 for a broad overview of how NPs have leveraged modality-specific components for different tasks.


---
Garnelo, Marta, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Ali Eslami. "Conditional neural processes." In International conference on machine learning, pp. 1704-1713. PMLR, 2018.

Garnelo, Marta, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J. Rezende, S. M. Eslami, and Yee Whye Teh. "Neural processes." arXiv preprint arXiv:1807.01622, 2018.

Saurav Jha, Dong Gong, Xuesong Wang, Richard E. Turner, Lina Yao. "The neural process family: Survey, applications and perspectives." arXiv preprint arXiv:2209.00517, 2023.

Reviewer comment (Presentation):
Table 3 might be improved by including "ours" in relevant rows for readability.

Author response:
Thank you for your suggestion. We have updated Table 3 to include "ours" in the relevant rows to improve readability.

Reviewer comment (Experiments):
As described in section 4.1, how much does increasing the number of iterations affect the performance of the model?

Author response:
Could you please clarify which iterations you are referring to?

Reviewer comment (Reproducibility):
What were the resources used for training, and what qualified as "too computationally expensive" for Non-Diagonal variants trained on CelebA?

Author response:
In the experiment, we used an Nvidia P100 (16GB) GPU. "Too computationally expensive" in this case refers to experiments where the model could not be trained due to memory usage.

Reviewer comment (Experiments):
The main weakness that I see in this work lies in the trade-off between constant memory usage and computational efficiency. To maintain constant memory, the input data is divided into batches, requiring multiple iterative updates rather than a single computation. This approach, while memory-efficient, could potentially lead to increased computational time, as the number of updates is inversely proportional to the batch size. The figures provided offer valuable insights into the relationship between block size and time within the CMABNP model itself. However, to fully substantiate the efficiency claims, it would be advantageous to expand the experimental setup to include comparisons of the CMABNP model with traditional update rules. Such an experiment could involve running the same model with and without the new update rule e.g. across CalebA32, 64, and 128. This would not only highlight the benefits of the proposed method in a more diverse context but also guide users in understanding the practical trade-offs when implementing this model in real-world scenarios.

Author response:
**Computational Efficiency of the update process:** Given the original input data $\mathcal{D}$ and a new set of datapoints $\mathcal{D}_{U}$ to update the model with, the original update process involves computing CMAB from scratch using $\mathcal{D} \cup \mathcal{D}_U$. In contrast, CMAB's update process leverages a rolling summation trick using only $\mathcal{D}_U$ to reduce the overall amount of computation. 

As a result, the amount of computation required for CMAB when using the original update process is $\mathcal{O}((|\mathcal{D}| + |\mathcal{D}_U|) |L_B| + |L_B|^2 + |L_B| |L_I| + |L_I|^2)$.

In contrast, CMAB's update process requires $\mathcal{O}(|\mathcal{D}_U| |L_B| + |L_B|^2 + |L_B| |L_I| + |L_I|^2)$ computation. Since $|L_I|$ and $|L_B|$ are pre-specified constants, this update process is constant computation per new datapoint.  

The computational efficiency gain by leveraging CMAB's update process is thus $\mathcal{O}(|\mathcal{D})| |L_B|)$, meaning that the efficiency gain scales with the size of the input data. This makes CMABs naturally extra useful in streaming data settings where the amount of contextual data increases over time.  

**Runtime Comparison:** We have included plots in the Appendix (Figure 11) which provide several comparisons of the runtimes. 

In Figure 11a, we compare CMANP using the efficient update process and the traditional update process, showing that the efficient update process is initially similar in runtime to the traditional update process. However, as the number of context datapoints increases (i.e., updates are performed) over time, the traditional update process' runtime grows linearly while our proposed efficient update process' runtime remains constant regardless of the number of tokens.

In Figure 11a, we also compare the runtime of the update process of CMAB-based NP (CMANPs) with Perceiver's iterative attention-based NP model (LBANPs) and a transformer-based NP model. We see that the CMAB-based model only requires a constant amount of time to perform the update. In contrast, Perceiver's iterative attention-based model's update runtime scales linearly and Transformer's update runtime scales quadratically.


In Figure 11b, we compare the querying (inference) runtime of CMANP with LBANPs (Perceiver's iterative attention-based model), TNPs (Transformer-based model). We see that CMANPs and LBANPs stay constant while the transformer-based model (TNP) scales quadratically in runtime.

Reviewer comment (Experiments):
While the paper posits the CMABNP model as particularly useful for streaming data settings like e.g. contextual bandits and Bayesian optimization, it lacks empirical evidence to support these claims. Given the relevance of such applications, it would strengthen the paper to include experiments that demonstrate the model's performance in these specific scenarios. Furthermore, a direct comparison with current state-of-the-art implementations, such as PFNs4BO [1], would not only validate the model's utility in practical settings but also offer a clearer understanding of its position relative to existing methods.

Author response:
A main point about streaming data setting is that the data is received batch by batch. As such, CMAB and CMANP are highly useful as they guarantee that the updates will be performed in constant computation and constant memory. We showed empirically their constant computation updates property in Figure 11 (Appendix) and constant memory property in Figure 3. In contrast, alternative attention methods would require computing from scratch, resulting in computation scaling linearly or quadratically. Furthermore, CMAB and CMANP do not require storing the previous context data when performing their updates, making it highly valuable in low-resource settings. Finally, in the appendix, we also included results of CMANPs applied to Contextual Bandits (a streaming data setting), showing CMANPs are competitive with prior state-of-the-art NPs. As such, we believe that a future application of CMABs (not limited to Neural Processes) to various streaming data settings would be an interesting research direction.

Reviewer comment (Reproducibility):
Lastly, the paper could be enhanced in the provision of implementation code.

Author response:
We will share the code alongside the camera-ready version. If it would help with the review process, we can prepare an anonymized version.

Reviewer comment (Evaluation):
Could you detail the computational efficiency of the update process and provide comparative runtimes for the models?

Author response:
**Computational Efficiency of the update process:** Given the original input data $\mathcal{D}$ and a new set of datapoints $\mathcal{D}_{U}$ to update the model with, the original update process involves computing CMAB from scratch using $\mathcal{D} \cup \mathcal{D}_U$. In contrast, CMAB's update process leverages a rolling summation trick using only $\mathcal{D}_U$ to reduce the overall amount of computation. 

As a result, the amount of computation required for CMAB when using the original update process is $\mathcal{O}((|\mathcal{D}| + |\mathcal{D}_U|) |L_B| + |L_B|^2 + |L_B| |L_I| + |L_I|^2)$.

In contrast, CMAB's update process requires $\mathcal{O}(|\mathcal{D}_U| |L_B| + |L_B|^2 + |L_B| |L_I| + |L_I|^2)$ computation. Since $|L_I|$ and $|L_B|$ are pre-specified constants, this update process is constant computation per new datapoint.  

The computational efficiency gain by leveraging CMAB's update process is thus $\mathcal{O}(|\mathcal{D})| |L_B|)$, meaning that the efficiency gain scales with the size of the input data. This makes CMABs naturally extra useful in streaming data settings where the amount of contextual data increases over time.  

**Runtime Comparison:** We have included plots in the Appendix (Figure 11) which provide several comparisons of the runtimes. 

In Figure 11a, we compare CMANP using the efficient update process and the traditional update process, showing that the efficient update process is initially similar in runtime to the traditional update process. However, as the number of context datapoints increases (i.e., updates are performed) over time, the traditional update process' runtime grows linearly while our proposed efficient update process' runtime remains constant regardless of the number of tokens.

In Figure 11a, we also compare the runtime of the update process of CMAB-based NP (CMANPs) with Perceiver's iterative attention-based NP model (LBANPs) and a transformer-based NP model. We see that the CMAB-based model only requires a constant amount of time to perform the update. In contrast, Perceiver's iterative attention-based model's update runtime scales linearly and Transformer's update runtime scales quadratically.


In Figure 11b, we compare the querying (inference) runtime of CMANP with LBANPs (Perceiver's iterative attention-based model), TNPs (Transformer-based model). We see that CMANPs and LBANPs stay constant while the transformer-based model (TNP) scales quadratically in runtime.

Reviewer comment (Experiments):
Is there empirical evidence to support the model's utility in streaming data settings?

Author response:
A main point about streaming data setting is that the data is received batch by batch. As such, CMAB and CMANP are highly useful as they guarantee that the updates will be performed in constant computation and constant memory. We showed empirically their constant computation updates property in Figure 11 (Appendix) and constant memory property in Figure 3. In contrast, alternative attention methods would require computing from scratch, resulting in computation scaling linearly or quadratically. Furthermore, CMAB and CMANP do not require storing the previous context data when performing their updates, making it highly valuable in low-resource settings. Finally, in the appendix, we also included results of CMANPs applied to Contextual Bandits (a streaming data setting), showing CMANPs are competitive with prior state-of-the-art NPs. As such, we believe that a future application of CMABs (not limited to Neural Processes) to various streaming data settings would be an interesting research direction.

Reviewer comment (Experiments):
Can you clarify whether parallel computation was used in the experiments? And if so, how does the model perform when processed sequentially, as would be common in low-resource environments?

Author response:
Our implementation is a sequential version. We detailed the parallel computation in the Appendix as potentially useful information for future implementations. We have clarified this in the updated draft.

Reviewer comment (Reproducibility):
Could you share the code?

Author response:
We will share the code alongside the camera-ready version. If it would help with the review process, we can prepare an anonymized version.