Multi-Timescale Dynamics Model Bayesian Optimization for Plasma Stabilization in Tokamaks

Thank you kindly for your detailed review. We have corrected all the typos you pointed out. Here are the answers to your questions:

Lessons: During the experiments, we realized the importance of having a compressed representation of the actuator space. The high data efficiency of our approach would not have been possible if the search space had been much more granular. At he same time, having a high-fidelity model was particularly important, as we saw it clearly outperforms models with more naive and less sophisticated priors. During our experiment, we also observed that the suggestions for ECH provided by humans sometimes differed from our approach. We think investigating ways to incorporate these preferences on the fly would be a very exciting avenue of research.

Using self-attention: We briefly experimented with an architecture that included self-attention while developing our model. Although we observed good predictive performance on the training data, the performance on the test data was poor, indicating overfitting. We speculate this is due to insufficient data and the overall tasks' difficulty (it has been shown that transformers sometimes only become good with a large amount of data (https://arxiv.org/abs/2106.04554)). For this reason, we stopped using transformers early during our research. Instead, we chose the RPNN architecture, which has been known to achieve state-of-the-art results on fusion data (https://arxiv.org/pdf/2404.12416). The paper cited looks into this question more deeply and analyzes GRU vs LSTM. Although we did not focus on this question in this work, we feel that designing appropriate attention-based architectures for nuclear fusion dynamics prediction is a very interesting research direction.

Thank you for reviewing our paper. If you feel that our responses to other reviewers improve the quality of the paper, we would be very thankful if you would consider changing your score accordingly.

Thank you kindly for your review. We have addressed your questions and comments in detail below.

Definition of frequencies: By frequencies, we mean the resolution of the input signal. The RPNN takes as input the full step-by-step actuator signals, which have a frequency of 50Hz, and makes predictions on a recursive, step-by-step basis, at the same frequency. By contrast, the Gaussian process model takes only the average ECH and betaN as input, which corresponds to a low-frequency smoothed version of the ground-truth actuator signals, and outputs the residual of time to tearing mode. (In other words, the RPNN is a high-resolution model, whereas the Gaussian process is a low-resolution model.) Based on this discussion, we have substituted the sentence you indicated with a clearer one.

Comparison with GP with marginal statistics: As you suggested, we have added the relevant comparison with a GP with marginal statistics as a prior mean. The table below shows the cumulative regret value after 500 iterations under the setup mentioned in the offline experiments setup. We observed that a prior mean corresponding to the historical mean of the data yields similar performance to that of a zero mean prior, indicating that a constant prior is insufficient for the present setting.

Initial conditions and context: The experiment configuration determines the initial conditions, which are agreed upon with DIII-D scientists before the experiment. The context for the GP is the target betaN, which is also decided before the experiment. We will make this clearer in the revised version.

Details of recurrent probabilistic neural network: Network Architecture -

Encoder :

1. FC layer (input dim x512)
2. FC layer (512x512)

Memory Unit :

1. GRU Block (512 x 256)

Decoder : Residual connections used between FC layers

1. FC Layer (256x512)
2. FC layer (512x512) repeated 8 times
3. FC layer (512x128)

This connects to two outputs

1. Mean head (128 x output dim)
2. Log Var head (128 x output dim)

The network predicts distribution parameters as output; hence, we train with a log likelihood loss. We use the Adam optimizer with a learning rate of 3e-4 and weight decay of 1e-3. We also use early stopping (patience = 250 epochs) on a validation set with 10% of the data points. We will include these details in the appendix of the revised manuscript.

Equation (5): We train a binary classifier that outputs a probability of tearing mode at time t given the state s_t and action a_t. The output specifies a Bernoulli distribution (conditioned on s_t and a_t), which we sample to predict the probability of a tearing mode occurring. We will make this clearer in the revised manuscript.

Experiments with varying lengthscale: The main goal of varying the lengthscale is to show that our method is robust to the choice kernel and hyperparameters, particularly compared to other approaches. We will state this more clearly in the revised manuscript.

Choice of acquisition function: We ran more offline experiments with different acquisition functions:

The numbers show the cumulative regret at the end of 500 iterations. We noticed similar performance across all acquisition functions. We chose UCB because it allows easy tuning of exploration vs exploitation during the actual experiment.

Lengthscale for Matern kernel: As mentioned above, the offline experiments with varying lengthscales aim to analyze our approach's robustness under varying model specifications. With this in mind, we first observed the lengthscale obtained by log-likelihood maximization over the full data set, then chose lengthscales for the offline experiments that varied from the log-likelihood maximum by at most an order of magnitude. For the live experiments, we used the lengthscales that yielded the highest fit according to the log-likelihood loss.

GP data efficiency: Thank you for pointing this out. We will revise the statement accordingly.