Dimension Agnostic Neural Processes

[Q3] Fine-Tuning Issues and Extrapolation

Following your suggestion, we conducted additional fine-tuning experiments on 5 d GP regression data to analyze the extrapolation ability of our method. In this experiment, we aim to compare not only the performance of a single DANP model against the baselines but also evaluate and compare multiple variants of DANP trained on different dimensional GP data. Specifically, we include DANP models trained on {1,2}, {3,4}, {2,4}, and {2,3,4} dimensional GP data, as well as the corresponding DANP models where sinusoidal PE is replaced with RoPE.

The results in Table R.14 clearly demonstrate that DANP outperforms the baselines in extrapolation few-shot scenarios, showcasing its robustness in handling these challenging tasks. Additionally, we observe that the DANP trained with 1,2d RoPE shows a notable improvement in generalization performance when provided with a few-shot setting. However, despite this improvement, its performance on the target data remains inferior compared to other DANP training settings, such as those utilizing higher-dimensional data ({3,4}, {2,4}, or {2,3,4}) or sinusoidal PE.

Table R.14 Comparison of fine-tuning performance between DANP with various settings and the baselines. Here, we use the few-shot 5d GP regression task with RBF kernel for the evaluation. We also compare the performances for both full finetuning and freeze finetuning for all models.

Table R.12 Comparison of zero-shot performance between DANP trained with the sinusoidal PE and the RoPE. Here, each method trained with {1, 2}d GP dataset with RBF kernel while performing inference on the {1, 2, 3, 4, 5}d GP dataset with RBF kernel.

Table R.13 Comparison of zero-shot performance between DANP trained with the sinusoidal PE and the RoPE. Here, each method trained with {3, 4}d GP dataset with RBF kernel while performing inference on the {1, 2, 3, 4, 5}d GP dataset with RBF kernel.

[Q2] GP Regression Setup and Zero-Shot Evaluation

Thank you for your insightful comment. The reason we set our experimental configurations to {2,4} and {2,3,4} was to comprehensively demonstrate our model’s performance in interpolation and both increasing and decreasing extrapolation scenarios.

Specifically, training on {2,4} was chosen to first show that the model works well on the interpolative dimension of 3. Additionally, we aimed to evaluate its ability to generalize to the decreasing extrapolation (e.g., 1-dimensional) case, as well as the increasing extrapolation to 5 dimensions and further extrapolation to 7 dimensions. Although some might consider 1-dimensional inference as learned due to our training on {2,4} dimensions, it remains an extrapolation task because our attention mechanism must update the representation space using only the first dimension’s positional encoding without higher-dimensional support.

While we could explore even larger extrapolative dimensions, GP data requires an exponentially growing number of context points as dimensions increase, making inference challenging. Therefore, 7 dimensions provided a balance for feasible inference. The subsequent training on {2,3,4} dimensions was included to show a clear performance gain on both interpolative and extrapolative tasks, underscoring how additional dimensions during training can boost performance across nearby dimensions.

However, we clearly understand the importance of analyzing our method's extrapolation capabilities as compared to its performance in interpolation settings. To address this, we followed your advice and conducted additional experiments by training on the {1,2} and {3,4} dimensional cases, then evaluating the results on {1,2,3,4,5} dimensional test data.

Here, we train DANP utilizing both sinusoidal PE and RoPE to further analyze their generalization ability. Tables R.12 and R.13 present the performance of DANP when trained on data from {1,2} dimensions and {3,4} dimensions, respectively.

From Table R.12, we observe that when trained on the limited range of {1,2} dimensions, both positional embedding methods fail to learn sufficient general features, leading to lower generalization performance compared to training on {2,4} or {2,3,4} dimensions. This result emphasizes the importance of training on higher-dimensional data to capture general features that enable better generalization to unseen dimensions. A similar pattern is evident in Table R.13.

However, a distinct trend emerges in Table R.12 compared to Tables R.9 and R.10. While both sinusoidal PE and RoPE performed similarly when sufficient general features could be learned from more diverse training dimensions, RoPE demonstrates noticeably weaker generalization ability than sinusoidal PE when the training data is limited to the narrow dimensional range of {1,2}. This result highlights the dependency of RoPE on richer training data which contains richer general features to achieve high generalization ability.

In the final manuscript, we will add this discussion in the main paper's conclusion and clearly reference and analyze the related experimental results from the appendix within the main experiment section to provide a more comprehensive analysis.

Table R.9 Comparison of zero-shot performance between DANP trained with the sinusoidal PE and the RoPE. Here, each method trained with {2, 4}d GP dataset with RBF kernel while performing inference on the {1, 2, 3, 4, 5}d GP dataset with RBF kernel.

Table R.10 Comparison of zero-shot performance between DANP trained with the sinusoidal PE and the RoPE. Here, each method trained with {2, 3, 4}d GP dataset with RBF kernel while performing inference on the {1, 2, 3, 4, 5}d GP dataset with RBF kernel.

Table R.11 Comparison of fine-tune performance between DANP trained with the sinusoidal PE and the RoPE. Here, each method trained with {2, 4}d GP dataset or {2, 3, 4}d GP dataset with RBF kernel while performing few-shot training on the 1d GP dataset with RBF kernel. Here, we report the performance for both the full finetuning and free finetuning.

[Q1] Positional Embedding Extrapolation

Thank you for the constructive and informative comment regarding the limitations of sinusoidal positional encoding. As you pointed out, many previous works have shown that sinusoidal positional encoding tends to perform poorly in terms of generalization when extrapolating to longer sequence lengths for Large Language Models. In response to this, approaches like RoPE have been proposed and used to address these limitations. While sinusoidal positional encoding successfully handled interpolation and extrapolation in our experimental settings, as you suggested, RoPE could potentially improve this performance. Therefore, we conducted additional experiments using a modified RoPE-based encoding tailored for the DAB module.

In our implementation, we retained the basic formulation of RoPE while ensuring different positional encodings for the x and y dimensions, similar to the approach we used with DAB. Specifically, we distinguished the embeddings added to queries and keys from x and y by alternating the cosine and sine multiplications for each. For example, if for x we calculate $q_{1x} \cdot \cos(\text{pos}) + q_{2x} \cdot \sin(\text{pos})$, then for y, we compute $q_{1y} \cdot \sin(\text{pos}) + q_{2y} \cdot \cos(\text{pos})$.

Using this modified positional encoding, we conduct additional experiments on the zero-shot and the fine-tune scenario in Gaussian Process regression tasks using the same settings in the main paper to evaluate RoPE’s impact on our model’s performance.

We conducted ablation experiments on sinusoidal PE and RoPE in a zero-shot scenario by inferring on 1D, 2D, 3D, 4D, and 5D GP regression data using DANP models trained on 2D and 4D GP regression data, as well as on 2D, 3D, and 4D GP regression data. The results, presented in Tables R.9 and R.10, indicate that while sinusoidal PE consistently outperforms RoPE in the 1D case, their performance is largely similar across other dimensions. This suggests that for these scenarios, both sinusoidal PE and RoPE exhibit comparable interpolation and extrapolation capabilities.

We also conducted experiments using the trained models to perform few-shot learning on 1D GP regression, following the setup in the main paper. As shown in Table R.11, while there were some performance differences in the zero-shot setting for the 1D GP regression task, these differences largely disappeared after few-shot fine-tuning. This indicates that the choice of positional embedding—whether sinusoidal PE or RoPE—has minimal impact on performance once the model is fine-tuned.

In the existing experimental settings presented in the main paper, we did not observe significant differences in performance or generalization ability between sinusoidal PE and RoPE. However, in experiments related to Q2, we identified certain scenarios where differences emerged. For further details and analysis, please refer to our response to Q2.

[Q5] Additional Metrics for Model Evaluation

Thank you for your question regarding additional model evaluation. In response to your suggestion, we measured and report the following 6 additional metrics: 1) Mean Absolute Error (MAE)**, 2) Root Mean Square Error (RMSE), 3) Coefficient of Determination ($R^2$), 4) Root Mean Square Calibration Error (RMSCE), 5) Mean Absolute Calibration Error (MACE), 6) Miscalibration Area (MA). Except for $R^2$, lower values for all these metrics indicate better alignment with the target and improved calibration performance. We conducted the evaluation using models trained on a 1d GP regression task, comparing our method with the baselines. The results, summarized in Table R.16, demonstrate that DANP achieves the best performance across a range of metrics. This observation reaffirms that DANP not only outperforms in terms of NLL but also achieves improved performance in calibration-related metrics compared to the baselines. These additional evaluations highlight the robustness of our method across diverse aspects of model performance. We will include this analysis in our revised manuscript to provide a more comprehensive evaluation of our approach.

Table R.16 Results on additional metrics containing MAE, RMSE, $R^2$, RMSCE, MACE, and MA on 1d GP regression task. Except for $R^2$, lower values for all these metrics indicate better alignment with target and improved calibration performance.

[Q6] Misc

Thank you for pointing out the errors in the code and tables. We appreciate your feedback, and we will make sure to address these issues. Specifically, we will ensure that the main paper accurately reflects the additional experimental results and provides a thorough analysis as per your suggestions. We are committed to refining the final manuscript accordingly and will make the necessary revisions to ensure clarity and consistency throughout.

[Q4] Video Completion Task

Thank you for suggesting the evaluation of zero-shot performance on the video completion task. We agree that reporting the zero-shot performance of DANP in this task provides a stronger contribution and maintains consistency with the GP regression experiments. Based on that actually, we already included zero-shot results in the main paper's Table 3 (b), indicated with a dagger symbol.

The results clearly demonstrate that even without fine-tuning, DANP achieves performance that surpasses the few-shot fine-tuning results of other models. This highlights DANP's robust generalization capabilities and further validates its effectiveness in handling complex tasks like video completion, aligning well with the strengths observed in our GP regression experiments.

Also, based on your suggestion, we conducted an additional experiment on the CelebA landmark task to further demonstrate the capabilities of our method. In the standard CelebA landmark task, the goal is to predict the locations of five facial landmarks: left eye, right eye, left mouth corner, right mouth corner, and nose, based on a single image. However, since Neural Processes predict a distribution over the target points using a given context, we adapted the CelebA landmark task to better fit this approach. We modified the task by combining the image's RGB values with the corresponding coordinates for each landmark, creating a 5-dimensional input. The output was restructured as a 5-dimensional label representing which of the five facial regions the prediction corresponds to. This setup allowed us to train and evaluate the model in a way that aligns with the predictive distribution framework of Neural Processes.

For the experiment, we used pre-trained models for the baselines, specifically the CelebA image completion models, while we trained DANP on both the EMNIST dataset and CelebA image completion tasks. This approach allowed us to assess the performance of DANP under a slightly modified but challenging setup, testing its ability to generalize across different types of tasks. Table R.15 validates that DANP still performs well on the different types of tasks compared to other baselines. And for the zero-shot scenario, DANP achieves 1.171±0.020 for the context dataset and 0.252±0.003 for the target dataset. These results demonstrate that although the target likelihood of zero-shot DANP is lower compared to that of fine-tuned baselines—primarily due to variations in both input and output dimensions from the training data—DANP quickly surpasses other baselines after fine-tuning. This highlights DANP's robust ability to generalize effectively in challenging zero-shot scenarios while rapidly improving with minimal fine-tuning.

Table R.15 Experimental results on the modified CelebA landmark task. Here, we fine-tuned baselines with 100-shot CelebA landmark dataset.

Thank you for reading our response and asking follow-up questions. Actually, yes, we rotated the hidden states of $x$ and $y$ for each $i$-th dimension to properly apply RoPE in the DAB module. And we'll make dagger more visible following your suggestion. Please feel free to let us know if you have any further questions or would like to discuss anything!

Thank you for your positive review and the constructive comments that suggested directions to better understand and improve our paper. We will structure the final manuscript carefully to make it more readable and self-contained.

[W5-2] Meanwhile it seems a lot of results in the context are nearly the same in scales in Table 2-3. Details about the number of shots are missing in experiments, further weakening the results.

Thank you for the feedback on the experimental results. To clarify, the similar likelihoods for the context in our experiments were due to our approach of modeling the standard deviation for all models' decoders' Gaussian distribution prediction as $0.1 + 0.9 \cdot \text{softplus}$. This was done to ensure fair comparisons between models. By modeling the standard deviation in this way, we prevent the performance from being biased by differences in the way the standard deviation is handled across models, ensuring that the output form is consistent for both DANP and all baseline models.

For the 1-dimensional output case, the performance of context point prediction reaching a value of 1.38 indicates a nearly perfect reconstruction. This result shows that most models were able to reconstruct the context points well in the given model setting. The similar values observed are consistent with other works, like in [3].

Furthermore, we have provided all the information about the context and target set sizes used in the experiments in Appendix A.3. The term "shot" that seemed confusing refers to the information needed when experimenting with unseen tasks, and we have made sure to include this information either in the main text or Appendix A.3 for clarity.

References

[1] Foong A, Bruinsma W, Gordon J, et al. Meta-learning stationary stochastic process prediction with convolutional neural processes[J]. Advances in Neural Information Processing Systems, 2020, 33: 8284-8295.

[2] Yann Dubois, Jonathan Gordon, and Andrew YK Foong, Neural Process Family. http://yanndubs.github.io/Neural-Process-Family/

[3] Hyungi Lee, Eunggu Yun, Giung Nam, Edwin Fong, and Juho Lee. Martingale posterior neural processes. In International Conference on Learning Representations (ICLR), 2023.

[W3-2] Disagree that NP is the earliest to address uncertainty as CNP can also achieve this in experiments

Thank you for the insightful comment. It seems there was a misunderstanding in the message. As you correctly pointed out, CNP models are indeed designed to capture data (aleatoric) uncertainty. However, what we intended to convey in that sentence is that the Neural Processes paper is the first to attempt capturing model (epistemic) uncertainty (uncertainty in the underlying function f) using a global latent path. We will make this clear in the revision.

[W3-3] It seems several works [2-9] are not discussed in the literature

Thank you for your constructive suggestion on improving our manuscript. Following your request, we will include a new, dedicated related work section in the appendix of the final manuscript. This section will provide a broader and more detailed discussion of related works to make the paper more accessible and comprehensible for readers. We appreciate your recommendation and believe this addition will enhance the overall quality and clarity of our work.

[W5-1] The computational complexity towards the modules and other ablations are required in examining the performance.

We appreciate your insightful question regarding the time complexity. We will analyze the time complexity compared to the Transformer Neural Processes(TNP) both in theoretically and practically.

First in theoretically, let us denote $B$, $N$, $d_x$, $d_y$, $d_r$, $L_d$, and $L_l$ denote the batch size, the number of data points (union of context and target), the dimension of input $x$, the dimension of output $y$, the representation dimension, the number of layers in the deterministic path, and the number of layers in the latent path, respectively. The additional computational cost for the DAB module is $O(BN(d_x+d_y)^2)d_r)$, and for the latent path, it is $O(L_lBN^2d_r)$. Since the computational cost for TNP is $O(L_dBN^2d_r)$, the overall computational cost of DANP can be expressed as $O((L_l+L_d)BN^2d_r)+O(BN(d_x+dy)^2d_r$. Generally, since $N >> (d_x+d_y)^2$ holds, the dominant term in the computational cost can be approximated as $O((L_l+L_d)BN^2d_r).

For the practical time cost, we measure the time cost to train 5000 steps for the GP regression tasks and image completion tasks for TNP and DANP. The results are shown in Table R.8.

Also, regarding the ablation study, as you suggested, we believe that conducting an ablation study to examine how the DAB module and latent path affect model predictions provides a deeper understanding of each module's role. Consequently, in the main paper, Section 5.4 under the paragraph titled The Role of the DAB Module and Latent Path, we analyzed each module's function through ablation experiments.

In summary of the results in Table 4, across various experimental settings, we observed that adding only the DAB module allows the model to handle data of varying dimensions, but it does not provide additional performance gains. In contrast, including the latent path results in performance improvements (indicating an enhanced capacity to learn shared features across tasks) but lacks the ability to manage varying dimensional data. These results demonstrate that the DAB module enables DANP to handle data with varying dimensions, while the latent path enhances the model’s ability to learn generally shared features across tasks of different dimensions.

Table R.8. Wall clock time evaluation for the TNP and DANP in various settings. Here, we utilize RTX 3090 GPU for the evaluation.

[W2] Did not see sufficient results to illustrate uncertainty when the dimension of the output is high

Thank you for the constructive comment. As you noted, demonstrating DANP's ability to learn and handle tasks with varying output dimensions concurrently is crucial, especially as it highlights the model's robustness in high-dimensional outputs. To show this, we conducted experiments on image and video completion tasks in the main paper Table 3, training on both EMNIST data (output dimension 1) and CelebA data (output dimension 3) simultaneously. These experiments verified that our model can effectively learn and infer across different output dimensions. Furthermore, in the video completion task, we empirically demonstrated that DANP is capable of handling and making inferences on video data (output dimension 3) with only minimal task-specific finetuning, unlike other models. This reinforces the versatility and strength of our model in diverse scenarios.

[W3-1] It has been demonstrated that the Eq (19) is not a valid ELBO. Hence Line 266-269 should be revised.

Thank you for the insightful comment. As you pointed out, and as highlighted in [1], the ELBO loss we used does not provide the exact ELBO for the $- \log p_\theta(y_t|x_t, D_c)$, because we use $q(z|D_c)$ instead of $p(z|D_c)$. More precisely, the Maximum Likelihood Loss is a biased estimator of $- \log p_\theta(y_t|x_t, D_c)$, and the ELBO we used is a lower bound of the same quantity [2]. Therefore, both losses still share the same issue, and the effectiveness of each loss depends on the model.

Typically, the maximum likelihood loss tends to exhibit larger variance compared to variational inference, so, given our model's need to handle multiple varying dimensional tasks simultaneously, we opted for variational inference to ensure stability. However, as you mentioned, it is worth experimenting with other loss functions. Therefore, we plan to include the results from training with the maximum likelihood loss as well.

We conducted ablation experiments on the ML loss and VI loss using DANP trained on 2 and 4d GP data, as well as DANP trained on 2d, 3d, and 4d GP data. These experiments were performed in a zero-shot scenario by inferring on 1, 2, 3, 4, and 5d GP regression data. The results, presented in Tables R.6 and R.7, show that while ML loss occasionally yields better log-likelihoods for context points, the VI loss consistently provides superior performance for the target points, which are of greater interest during inference. This trend is particularly evident in experiments trained on 2, 3, and 4d GP data. These findings demonstrate that using the VI loss for training DANP is generally more beneficial for improving generalization compared to the ML loss.

Table R.6 Comparison of zero-shot performance between DANP trained with the variational loss and the maximum likelihood loss. Here, each method trained with {2, 4}d GP dataset with RBF kernel while performing inference on the {1, 2, 3, 4, 5}d GP dataset with RBF kernel.

Table R.7 Comparison of zero-shot performance between DANP trained with the variational loss and the maximum likelihood loss. Here, each method trained with {2, 3, 4}d GP dataset with RBF kernel while performing inference on the {1, 2, 3, 4, 5}d GP dataset with RBF kernel.

[W4] The zero-shot and fine-tune experiments are not appropriate in evaluation as NP families require the context and amortize the few-shot adaptation without gradient updates. I am not convinced by the meaning of fine-tune or zero-shot in the task concept

Thank you for the constructive comment. Although this is not the reviewer utkY's first question, we would like to address this point first to ensure there is no misunderstanding about the zero-shot and fine-tuning scenarios before moving on to discuss other questions. It seems there was a misunderstanding regarding the experimental setup in our paper. By "zero-shot" and "fine-tuning" in our experiments, we did not mean performing zero-shot inference without a context set or fine-tuning on a given context before inferring on the target set. Instead, we refer to the process where a model trained on tasks consisting of a context set and target set is then used to perform inference on unseen tasks, which involve context and target sets with input dimensions that were not part of the original training data.

For example, in Table 2, the model trained on tasks sampled from 2d, 3d, and 4d dimensional GPs was used to infer on unseen dimensional tasks, such as 1d, 5d, and 7d dimensional GPs. These unseen tasks still contain context and target sets, and inference is performed based on the given context. In this context, "zero-shot" refers to directly using DANP to infer on unseen dimensional tasks without additional training, which is a capability that other models cannot achieve. On the other hand, "fine-tuning" refers to experiments where a small number of training samples from the unseen task are used to further train both the baseline models and DANP to observe how performance improves.

[W1] Lacks theoretical analysis

While our proposed DANP model lacks theoretical analysis, we have provided extensive empirical analysis demonstrating that DANP performs more robustly in zero-shot and few-shot scenarios on varying-dimensional tasks compared to other baseline models. Specifically, in the GP regression task, we show that DANP achieves performance comparable to that of fully trained baselines solely through zero-shot inference, even without being trained on 1d regression data. This result demonstrates DANP's strong capability to generalize and capture shared features across previously unseen tasks with different dimensions.

However, as you suggested, incorporating theoretical analysis would strengthen our paper further, and we see it as a valuable direction for future work.

Thank you for your prompt response.

You may know that we are conducting in-context learning at the task level but not the data point level in our paper and [1]. This means that, without updating the weights, the Neural Process utilizes its pre-amortized knowledge to infer a new given task. In the ICL setting, inference may initially be inaccurate for the new task when relying solely on previously learned knowledge. However, by providing examples from new tasks, the model adjusts its inference, enabling more accurate predictions.

In our paper, we conducted purely zero-shot experiments at the task level to demonstrate that when the model learns tasks of a similar nature (for example different dimensional GP tasks), it can still achieve strong performance even in purely zero-shot inference at the task level without additional examples, unlike the ICL setting.

We are curious to hear your thoughts on this experimental setup and empirical validation. Additionally, we would like to understand which aspects of the Neural Process framework you believe this approach might conflict with.

Thanks for the feedback. I ever read the mentioned paper "In-Context In-Context Learning with Transformer Neural Processes. " However, I'm afraid this work has some misunderstanding about the zero-shot learning or testing.

Actually, in-context learning is also kind of few-shot learning, which depends on the context information. This is not typical zero-shot testing.

Thank you for your response to our rebuttal. However, we are curious why you believe the zero-shot setting contradicts the original setting of Neural Processes. Neural Processes are fundamentally designed to learn from various tasks and amortize the acquired knowledge to solve new, unseen tasks.

Aligned with this goal, while different from our paper they do not focus on creating models agnostic to unseen dimensions, studies like [1] extend Transformer Neural Processes for in-context learning, solving new tasks with the same dimensions (e.g., training on CelebA and testing on CIFAR10 image completion tasks). In this case, these models also tackle tasks in a zero-shot manner without fine-tuning on new tasks, instead relying on in-context learning with a few data points.

Thus, we are keen to understand which aspects of the Neural Process framework you believe are unsuitable for such experiments and research directions.

References

[1] Ashman, M., Diaconu, C., Weller, A., & Turner, R. E. (2024). In-Context In-Context Learning with Transformer Neural Processes. arXiv preprint arXiv:2406.13493.

Thank you for your dedication and interest in our paper. As the author and reviewer discussion period approaches its end, we are curious to know your thoughts on our rebuttal and whether you have any additional questions.

[W2] The paper should include an analysis or discussion on how the components of DANP contribute to its predictions, especially given its complex architecture involving the DAB and Transformer-based latent path.

Thank you for the insightful comment. As you suggested, we believe that conducting an ablation study to examine how the DAB module and latent path affect model predictions provides a deeper understanding of each module's role. Consequently, in the main paper, Section 5.4 under the paragraph titled The Role of the DAB Module and Latent Path, we analyzed each module's function through ablation experiments.

In summary, across various experimental settings, we observed that adding only the DAB module allows the model to handle data of varying dimensions, but it does not provide additional performance gains. In contrast, including the latent path results in performance improvements (indicating an enhanced capacity to learn shared features across tasks) but lacks the ability to manage varying dimensional data. These results demonstrate that the DAB module enables DANP to handle data with varying dimensions, while the latent path enhances the model’s ability to learn generally shared features across tasks of different dimensions.

[Q2] Are there any plans to conduct longitudinal studies to evaluate the long-term performance and stability of DANP in dynamic environments?

Thank you for inquiring about our future work plans. As demonstrated in our experiments with the MIMIC-III dataset, our ultimate goal is to develop a general foundation regressor model capable of handling a wide range of data structures and realistic scenarios, such as diverse time series data and cases with missing features. We view the DANP research as the initial step toward achieving this ambitious objective.

A key focus of our future work will be to extend the model’s ability to appropriately process inputs with varying dimensions, numbers of context and target points, and diverse data structures (for example there can be lots of different tasks with the same dimensional inputs such as EMNIST image completion and 2d GP regression task). Developing a model that can flexibly adapt to such variability without specific data processing based on inductive bias while providing accurate and reliable inferences across these scenarios remains a critical challenge and an exciting direction for further exploration.

[W1, Q1] It could be benefit from additional experiments or a theoretical discussion on how DANP might perform in non-regression tasks. Are there any modifications needed for effective application in these domains?

Following the reviewer's insightful suggestion, we conducted additional experiments on time series data using the blood pressure estimation task from the MIMIC-III dataset. Specifically, we assumed real-world scenarios where certain features from patient data might be missing, or entirely different sets of features could be collected. Under this assumption, we trained the model using only a subset of features from the MIMIC-III dataset and evaluated its performance when additional or different sets of features became available.

Specifically, we considered five features: T, Heart Rate, Respiratory Rate, SpO2, and Temperature. For pre-training, we utilized T and Heart Rate features, while for the fine-tuning scenario, we assumed only Respiratory Rate, SpO2, and Temperature features were available (this scenario can happen if we assume that we trained our model with data from hospital A and want to evaluate on the data in hospital B). We pre-trained the models with 32,000 training samples and fine-tuned them with only 320 samples. And here, we considered observations from time $0,...,t$ as context points and $t+1,...,T$ as target points. As shown in Table R.5, our DANP achieved strong performance in the time series blood pressure estimation task, demonstrating robustness and adaptability in this real-world scenario. These results are consistent with the findings presented in the main paper, further validating DANP's effectiveness in handling diverse and practical challenges.

And also our model can, of course, be applied to classification problems as well. However, since output predictions in classification are typically better modeled with likelihoods other than a Gaussian distribution, it would be beneficial to adapt the decoder model structure accordingly. By redesigning the decoder, our approach can effectively handle classification tasks.

Table R.5. Empirical results on the time series blood pressure estimation task.

Thank you for your dedication and interest in our paper. As the author and reviewer discussion period approaches its end, we are curious to know your thoughts on our rebuttal and whether you have any additional questions.

Thank you for the positive comments on our paper. We promise to properly add the additional experiments and discussions conducted in the discussion period to the final manuscript.

[W1] Design is complex, making replication and understanding challenging

We appreciate the valuable feedback aimed at making our paper clearer and more accessible for readers. Based on the discussions during the rebuttal period, we plan to revise the content in the final manuscript to improve readability and clarity, ensuring it is organized in a way that facilitates better understanding.

[W2] There’s no discussion of the model’s time and resource requirements, impacting assessments for practical use.

We appreciate your insightful question regarding the time complexity. We will analyze the time complexity compared to the Transformer Neural Processes(TNP) both in theoretically and practically.

First in theoretically, let us denote $B$, $N$, $d_x$, $d_y$, $d_r$, $L_d$, and $L_l$ denote the batch size, the number of data points (union of context and target), the dimension of input $x$, the dimension of output $y$, the representation dimension, the number of layers in the deterministic path, and the number of layers in the latent path, respectively. The additional computational cost for the DAB module is $O(BN(d_x+d_y)^2)d_r)$, and for the latent path, it is $O(L_lBN^2d_r)$. Since the computational cost for TNP is $O(L_dBN^2d_r)$, the overall computational cost of DANP can be expressed as $O((L_l+L_d)BN^2d_r)+O(BN(d_x+dy)^2d_r$. Generally, since $N >> (d_x+d_y)^2$ holds, the dominant term in the computational cost can be approximated as $O((L_l+L_d)BN^2d_r).

For the practical time cost, we measure the time cost to train 5000 steps for the GP regression tasks and image completion tasks for TNP and DANP. The results are shown in Table R.3.

Table R.3. Wall clock time evaluation for the TNP and DANP in various settings. Here, we utilize RTX 3090 GPU for the evaluation.

[W3, Q1] It would be helpful if the authors could provide more concrete examples and relate them to more complex, real-world applications

We appreciate your constructive suggestion regarding demonstrating DANP's applicability in additional practical scenarios in real-world applications. While the hyperparameter optimization task explored in our main paper is indeed a critical and highly relevant application, we agree that showcasing DANP's performance in other practical settings could further strengthen the contributions of our work. To validate the practicality, we conducted additional experiments on time series data using the blood pressure estimation task from the MIMIC-III dataset. Specifically, we assumed real-world scenarios where certain features from patient data might be missing, or entirely different sets of features could be collected. Under this assumption, we trained the model using only a subset of features from the MIMIC-III dataset and evaluated its performance when additional or different sets of features became available.

Specifically, we considered five features: T, Heart Rate, Respiratory Rate, SpO2, and Temperature. For pre-training, we utilized T and Heart Rate features, while for the fine-tuning scenario, we assumed only Respiratory Rate, SpO2, and Temperature features were available (this scenario can happen if we assume that we trained our model with data from hospital A and want to evaluate on the data in hospital B). We pre-trained the models with 32,000 training samples and fine-tuned them with only 320 samples. And here, we considered observations from time $0,...,t$ as context points and $t+1,...,T$ as target points. As shown in Table R.4, our DANP achieved strong performance in the time series blood pressure estimation task, demonstrating robustness and adaptability in this real-world scenario. These results are consistent with the findings presented in the main paper, further validating DANP's effectiveness in handling diverse and practical challenges.

Table R.4. Empirical results on the time series blood pressure estimation task.

Thank you for your dedication and interest in our paper. As the author and reviewer discussion period approaches its end, we are curious to know your thoughts on our rebuttal and whether you have any additional questions.

[W2, Q] The performance seems marginal compared to TNP in some cases. Share the confidence interval

Thank you for your comment regarding the performance and the confidence interval. We would like to first point out that our contribution is not about doing better than TNP for individual tasks, but showing that DANP is capable of processing variable dimension data seamleassely. Still, In response to your comment, while the performance figures in Table 1 (from-scratch GP regression) and Table 3 (image completion) may appear similar between DANP and TNP, it is important to note that DANP achieves differences in the second decimal place with the third decimal place standard deviations, which can be considered a significant improvement. More critically, we want to highlight that the DANP demonstrates substantial advantages in scenarios where zero-shot inference on new dimensional data or few-shot learning is required before inference. These strengths are clearly reflected in results such as Table 2, Table 3 (b), and the realistic Bayesian Optimization benchmark experiments shown in Figure 4.

We also agree with the reviewer's suggestion that additional evaluation through various metrics could further validate the effectiveness of DANP. As such, we have conducted further analyses with additional metrics, which are detailed in Appendix B.2. To further compare DANP with other baseline models using alternative metrics, we provide results for two additional metrics: the Continuous Ranked Probability Score (CRPS) and Empirical Confidence Interval Coverage. These results can be found in Table 13 in Appendix B.2. In summary, for a target confidence interval, DANP demonstrates relatively narrower intervals, delivering the best results among the models. Moreover, when comparing CRPS scores, DANP achieves better scores than the other models.

Also, we additionally measure some other metrics related to calibration. We measured and report the following 6 additional metrics: 1) Mean Absolute Error (MAE)**, 2) Root Mean Square Error (RMSE), 3) Coefficient of Determination ($R^2$), 4) Root Mean Square Calibration Error (RMSCE), 5) Mean Absolute Calibration Error (MACE), 6) Miscalibration Area (MA). Except for $R^2$, lower values for all these metrics indicate better alignment with the target and improved calibration performance. We conducted the evaluation using models trained on a 1d GP regression task, comparing our method with the baselines. The results, summarized in Table R.2, demonstrate that DANP achieves the best performance across a range of metrics. This observation reaffirms that DANP not only outperforms in terms of NLL but also achieves improved performance in calibration-related metrics compared to the baselines. These additional evaluations highlight the robustness of our method across diverse aspects of model performance. We will include this analysis in our revised manuscript to provide a more comprehensive evaluation of our approach.

Table R.2 Results on additional metrics containing MAE, RMSE, $R^2$, RMSCE, MACE, and MA on 1d GP regression task. Except for $R^2$, lower values for all these metrics indicate better alignment with target and improved calibration performance.

[W1] I find that more on the actual task considered here can help to appreciate more the significance and the benefits of this approach. In particular related how the GP regression and the image completion tasks benchmarks help to validate the broad applicability of the approach?

We appreciate your constructive suggestion regarding demonstrating DANP's applicability in additional practical scenarios in real-world application. While the hyperparameter optimization task explored in our main paper is indeed a critical and highly relevant application, we agree that showcasing DANP's performance in other practical settings could further strengthen the contributions of our work. To validate the practicality, we conducted additional experiments on time series data using the blood pressure estimation task from the MIMIC-III dataset. Specifically, we assumed real-world scenarios where certain features from patient data might be missing, or entirely different sets of features could be collected. Under this assumption, we trained the model using only a subset of features from the MIMIC-III dataset and evaluated its performance when additional or different sets of features became available.

Specifically, we considered five features: T, Heart Rate, Respiratory Rate, SpO2, and Temperature. For pre-training, we utilized T and Heart Rate features, while for the fine-tuning scenario, we assumed only Respiratory Rate, SpO2, and Temperature features were available (this scenario can happen if we assume that we trained our model with data from hospital A and want to evaluate on the data in hospital B). We pre-trained the models with 32,000 training samples and fine-tuned them with only 320 samples. And here, we considered observations from time $0,...,t$ as context points and $t+1,...,T$ as target points. As shown in Table R.1, our DANP achieved strong performance in the time series blood pressure estimation task, demonstrating robustness and adaptability in this real-world scenario. These results are consistent with the findings presented in the main paper, further validating DANP's effectiveness in handling diverse and practical challenges.

Table R.1. Empirical results on the time series blood pressure estimation task.

Thank you for your dedication and interest in our paper. As the author and reviewer discussion period approaches its end, we are curious to know your thoughts on our rebuttal and whether you have any additional questions.

Thank you for the positive comments on our paper. We promise to properly add the additional experiments and discussions conducted in the discussion period to the final manuscript.