We thank the reviewer for their time in evaluating our submission. While we note the critical tone of the review, we appreciate the opportunity to respond and clarify aspects of our work. We have made a sincere effort to address all points raised. All comments and questions will be addressed individually in an inline fashion.

Analytic likelihood: Do I understand it correctly that we are using reverse KL for learning and hence need an analytic likelihood?

The paper does indeed use a reverse KL loss for learning as mentioned in equation 2. The implementation requires a target density that can be evaluated up to a normalizing constant and is differentiable (usually using standard autograd tooling built into widely accepted libraries like PyTorch, TensorFlow or JAX). In the Bayesian setting, where the target is proportional to the product of the likelihood and prior, this requirement translates to only needing the likelihood to be differentiable and evaluable up to normalization. The likelihood does not need to have a closed‐form density; in fact, this is a motivating problem behind transdimensional inference, where the normalizing constant of each model is intractable. We will aim to clarify the text on this point.

Flow choice: I consider the use of autoregressive flows as a weakness... Why is that necessary? Could we use other flow architectures or do we need the specific autoregressive structure for the transdimensional part?

In the variational inference setting, inverse autoregressive flows are preferred over coupling flows due to:

1. Faster sampling during training (Kingma et al. 2017; Durkan et al. 2019). Autoregressive transformations use MADE which has an asymmetrical computational cost depending on the direction of the transformation. Inverse Autoregressive Flows are implemented to have computational cost $\mathcal{O}(1)$ (and are parallelizable) for sampling and evaluation of those samples as a by-product, which is the only operation used in the variational inference optimization loop. This means conversely the evaluation of out-of-distribution data points has the sequential cost $\mathcal{O}(d_{\mathrm{max}})$ due to sequential inverse evaluation operation that is often cited as the drawback of these flows for generative models. Since the use case is primarily to search over a large model space in a variational inference setting, and to obtain expectations with respect to the amortized variational density, this is not a limitation in our setting.

2. Higher expressiveness overall (Coccaro et al. 2024).

The desired factorization that is proved in Proposition 2.2 does rely on using the autoregressive property of the flow, however we do not see a barrier to deriving a CoSMIC architecture specifically for coupling flows in applications outside variational inference, such as simulation-based inference. This would be an interesting topic for future research. We will add some comments along these lines in the Discussion.

Jointly learning densities for models: You say that rather than training a density for each individual model separately, it is preferable to construct a single density. In which circumstances specifically is this likely to be true?

This approach is preferable in large model‐space settings. For instance, the model space of directed acyclic graphs (DAGs) scales super‐-exponentially: in the non‐linear DAG example, the model space of 10‐node DAGs has cardinality $|\mathcal{M}|=4 175098976430598143$. In such scenarios, it is infeasible to store individual normalizing flows for each model. We will revise our statement to include this explicit example.

Entropy gap + architectural limitations: How good is the entropy achieved in Experiment 1?

Figure 2 offers a holistic assessment of inference quality relative to a sampling baseline using RJMCMC. The entropy is reduced as flow expressivity increases; it also reduces with the size of the model probability (Figures 2, 5-10). The entropy gap scales with the model‐space size (Figure 4) when the flow architecture is held fixed. This matches expectations about neural density expressivity and indicates that flow design should be tailored to the problem domain. We will improve our description of this in the Experiments by bringing Figure 11 plus discussion to the main paper as a visual indicator of how good it is, and clarify the limits of the flow architecture in the Discussion.

Simplicity of benchmarks: They ... evaluate it on some relatively simple benchmarks.

We agree that the models used are initially relatively simple. That said, the robust variable selection problem is designed to be easily scaled to increased complexity in dimension and model space cardinality (see Figure 4, Appendix E.1). We also tackle non-linear DAGs, which are NP‐hard and non‐trivial in both computational complexity and setup. Our goal is not to dominate state‐of‐the‐art benchmarks but to introduce a novel variational method generally applicable to challenging domains with few comparable approaches. We will include a statement along these lines in the revision.

Scalability: What are potential scalability issues to high dimensional model/parameter spaces?

Thank you for raising this question. Our response to reviewer cEAN on scalability provides some background on potential scalability issues to high dimensional parameter spaces and high cardinality model spaces. In addition, in Section 3 we derive scaling properties for the GP-surrogate model distribution. In particular, high cardinality model spaces do require careful construction of a context encoder $\delta(m)$. This encoder is learned along with the flow, and is intended to non-linearly project $m$ to a higher dimension to improve expressiveness of the residual block described in Appendix A.3. We will add a paragraph to the discussion and experiment setup describing this consideration in the revised version, and refer the reader to the appendices for details.

Model index surrogate: $A GP$ seems ... unnecessarily technical vs just using a neural net

Thank you for this astute observation. We agree that in practice it is safe to default to using Monte Carlo gradient estimation using a categorical distribution in low cardinality problems, or a neural net like MADE in high cardinality problems. However, the GP approach offers a different framework with its own flexibility, offers competitive performance, and has the benefit of theoretical convergence properties that we provide in Section 3. We will add a summary of this clarification to the discussion in the revised paper.

Amortisation: In which sense is the approach “amortized” (clarification)?

Here we mean that, as discussed in Section 3, we train one shared recognition network $f_{\phi}$ that (i) is conditioned on model index $m$ to output $q_{\phi}(\theta_{m},u_{\setminus m}\mid m)$ in a single forward pass, and (ii) undergoes heavy optimization during training so that any subsequent inference over the posterior/target support requires only relatively cheap evaluation of the model distribution and evaluation of the IAF (no per-model VI runs). We will improve the text in Section 3 to make this clearer.

Additional References

Coccaro, Andrea, Marco Letizia, Humberto Reyes-González, and Riccardo Torre. 2024. “Comparison of Affine and Rational Quadratic Spline Coupling and Autoregressive Flows Through Robust Statistical Tests.” Symmetry 16 (8). https://doi.org/10.3390/sym16080942.

We thank the reviewer for their thoughtful and constructive comments, and we appreciate their interest in our paper. The feedback provided has helped us refine our approach, and we hope that our responses clarify and strengthen the manuscript. All comments and questions will be addressed individually in an inline fashion.

Presentation and clarity: Presentation is dense and paper is difficult to read. The paper should open up to a broader audience.

On reflection we agree, and will aim to use the additional page for the camera-ready version to improve overall paper accessibility and readability for a wider audience. This will be in addition to the specific changes discussed below.

Model weights variants: The point of 3 alternative solutions is lost when the evaluation largely skips this question. The performance of the MC gradient estimator as depending on model space is also not validated in detail.

Thank you for this insight. We will a) improve the text in Section 3 to make clearer the potential value and need for using different model weights estimators, and b) explicitly state in all experiments which particular estimator was used. Here, every experiment uses the Monte Carlo gradient estimator, whether as the main approach (Figures 2 & 3, Table 1) or in comparison to the variants (Figure 4). We discuss the MCG estimator performance in a separate comment below.

Detail of experimentation: The presentation in the main paper is a bit superficial and we do not learn very much about how the method really works.

Thank you – we agree. In line with the above responses we will use the additional manuscript space to expand upon experimental details and content. Some of this will entail bringing some of the content currently in the Appendix into the main paper (e.g. Figure 4 plus discussion, and if possible, Figure 11 and discussion). It will also involve improving the quality of the discussion, and bringing missing detail regarding the experimental setup (such as the true $\sigma$) into the main paper.

High cardinality MC estimator: How well the Monte Carlo estimator for high cardinality works in practice? If used for problem of low cardinality, is it roughly as good as the GP solution or substantially worse? If worse, should we expect it to be accurate for high cardinality?

In Appendix E.1, Figure 4, we conducted a comparison between the diagonal surrogate, categorical Monte Carlo gradient, and neural (MADE, Germain et al. 2015) Monte Carlo gradient (MCG) approaches on problems of increasing model space cardinality on two different levels of misspecification for the robust variable selection problem type. We fixed the number of iterations and the architecture of the normalizing flow (RQS). From this, we observed that the MCG methods are competitive with the Gaussian surrogate approach, and in some cases performed slightly better (though this is subject to tuning of both approaches). The benefits and drawbacks of each approach cannot be made in a broad statement, but rather are problem specific. For instance, we provided theoretical guarantees for the surrogate approach in Section 3. Also, there could be scope for future research into exploiting model space structure for the surrogate approach. The benefit of the MCG methods is, as you summarised, applicability to very high cardinality model spaces where we use a neural (MADE) architecture for the model distribution. Figure 4 shows comparable empirical performance on the robust variable selection problem in model space cardinalities up to $2^{24}$. The drawbacks are not concrete: we do not show the same theoretical guarantees, but this could be a topic of future research.

In order to communicate this comparison more clearly, we will move this content (Figure 4 plus discussion) to the main paper, and improve the clarity and quality of the discussion.

UCB convergence: Corollary 3.2 and the paragraph below that talks about convergence of the UCB-based solution but makes no reference to $\beta$, implying that already the mean (using $\beta\approx 0$) would work as well. Can you clarify what is happening here? I probably missed something about the claim.

Indeed, the value of $\beta$ should not affect the asymptotic convergence rate, according to our theoretical analysis, though, in practice, it helps the algorithm in the short term. This independence on $\beta$ is due to the result in Corollary 3.2 mainly only depending on maintaining a positive sampling probability over all models in the model space, which then allows us to apply the second Borel-Cantelli lemma to bound the posterior variance of the GP surrogate (Lemma B.2) and control the KL divergence of the approximation.

We sincerely thank the reviewer for their detailed and insightful comments, as well as for their strong interest in our work and methodology. Their expertise is evident in the depth of the review, and we are grateful for the opportunity to respond. We have carefully considered each of the points raised and have done our best to address all concerns. All comments and questions will be addressed individually in an inline fashion.

Experimental detail: Empirical analysis could be more comprehensive.

We agree, and we will use the additional page for the camera-ready version to expand upon experimental details and content. Some of this will entail bringing some of the content currently in the Appendix into the main paper (e.g. Figure 4 plus discussion, and if possible, Figure 11 plus discussion). It will also involve improving the quality of the discussion (emphasising important points), and bringing missing detail regarding the experimental setup (currently in Table 2) to enable more tangible understanding and discussion of the most relevant points.

Sparsity: If we’re interested in sparse modelling solutions, do we really need to increase the modelling complexity to $d_{max}$? Can we elaborate on introducing any "sparsity priors" to the proposed solutions?

Thank you for this insightful question. In terms of what is implemented in the experiments: We use a prior that induces “extra” sparsity in the real-data example for non-linear DAGs (Sachs et al 2005). The use of this prior is to further down-weight the probability of graphs with more edges in order to reach an acceptable level of closeness to the “consensus” graph in Sachs et al 2005. In no other experiment do we use sparsity-inducing priors.

As to the wider question: Introducing sparsity-inducing priors will tend to produce a posterior with little weight on non-sparse solutions (unless justified by the data). Because we know (see experiments and Discussion section) that most of the effort in approximating the posterior focuses on high posterior model probability models, this means that even though the flows are defined in $d_{max}$-dimensional space, the extra dimensions used in the higher-dimensional but lower-probability models will have little impact on minimising the loss. From this perspective the ability of VTI to approximate the posterior is not sensitive to this extra dimensionality, even though in practice the CoSMIC setup requires its specification. In general though, this question does open the door to future research on how to reduce these practical requirements (beyond hard-specification of a reduced $d_{max}$), perhaps via a more dynamic (run time) setup, though this may present implementational difficulties. Thank you for raising this question – we will include a discussion on this point in the Discussion of the revised version.

Figure 2: I did not completely understand Figure 2’s bottom row. Is there a discernible pattern when looking at the plots from left to right? I would expect the cross-entropy to decrease as we increase the complexity of the variational density.

We apologise for the lack of clarity in Figure 2. You are correct to expect the cross-entropy to decrease when looking from left to right. In Figure 2 there are actually two problem types: medium misspecification (panes 1, 2, 3), and high misspecification (panes 4, 5, 6). Left to right (increasing flow expressivity) for each problem type we observe decreasing $H$. We will improve the plot with a clear dividing line between the two problem types to improve readability, and also improve the caption and surrounding text to make this more clear.

DAG experiment baselines: Authors should discuss the relatively restricted baseline choice (DAGMA) in further detail.

Thank you, we agree and will include the cited papers in the discussion. So far, we have run additional benchmarks using DiBS/DiBS+ (Lorch et al. 2021) in addition to DAGMA for the 10-node MLP DAG synthetic study (please see the revised figures in the tables below, these have been added to Figure 3 in the revised paper), and we expect to run further benchmarks before the end of the reviewing period. We expect the additional benchmark performance to be close to VTI performance. (As an aside: we have found a small error in our plotting code, and DAGMA performance is now much closer to VTI in line with expectations.) What we are aiming to demonstrate with this analysis is that the performance of a general methodology like VTI can be competitive with an application-specific methodology (like DiBS, DAGMA). In general, of course, one would expect the application-specific methodology to perform best. We will improve the discussion on this point in the Experiments section, along with the new benchmark results.

Proposition 3.1: $This$ is relatively straightforward ... The paper would be strengthened if the space had been used for a discussion ...

We agree, and will recover this space to be used for the improvements discussed above/below.

MCMC acceptance probability: Also, when discussing reversible jump MCMC (Section E.4), the defined ratio for the death move is always greater than or equal to one. I believe should refer to the ratio within the minimum operator. Is that right?

Thank you for picking up on this error. We mistakenly wrote that the acceptance probability is the reciprocal of $\alpha_{\mathrm{birth}}$, whereas it should read

$r_{\mathrm{birth}} = \dfrac{p(y | \beta, m')p(\beta | m') p(m')}{p(y| \beta, m) p(\beta | m) p(m)}$,

$\alpha_{\mathrm{birth}} = \mathrm{min}$1,r_\mathrm{birth}$$,

\alpha_{\mathrm{death}} = \mathrm{min}\left1,\dfrac{1}{r_{\mathrm{birth}}}\right$$.

We will correct this in the manuscript. The acceptance probability is implemented correctly in the code.

Extra baseline results for figure 3

F1

SHD

Brier

AUROC

We thank the reviewer for taking the time to engage with our manuscript and for their feedback. We appreciate the effort to provide a thorough assessment of the work, and we value the opportunity to improve our submission. All comments and questions will be addressed individually in an inline fashion.

Baseline comparisons: Including comparisons with additional relevant baselines ... would better contextualize the method’s performance.

We agree that baselines must be appropriate to the developed methods. The goal of this paper is to introduce the very first method to achieve a flexible, trans-dimensional variational algorithm (VTI), which we obtain by developing the CoSMIC architecture for an inverse autoregressive flow. In this sense, the relevant baseline is the ground-truth, to show that the method both works and performs as expected. This is the case with the first experiment (Bayesian misspecified robust variable selection) in which we compare to the ground truth obtained by RJMCMC output. To our knowledge, there are no other baseline methods capable of producing a variational approximation to the transdimensional posterior.

For the second experiment (DAGs), we agree that additional baselines are more important. Here, because of the immense size of the model space a ground truth is difficult to obtain, and so we evaluated VTI performance via typical DAG diagnostics in Figure 3 and Table 1. We agree that in this setting comparison to a single method (DAGMA) is not sufficient. So far, we have run the additional benchmarks DiBS/DiBS+ (Lorch et al. 2021) (please see the revised results in the tables in response to zSgg). We expect to run further benchmarks before the end of the reviewing period. We expect the additional benchmark performance to be within the range of DiBS/DIBS+ and DAGMA. What we are aiming to demonstrate with this analysis is that the performance of a generic methodology like VTI can be competitive with an application-specific methodology. One would expect the application-specific methodology to perform best. We will improve the discussion on this point in the Experiments section. The additional results will be added to Figure 3.

We will use the additional page for the camera-ready version to expand upon experimental details and content. Some of this will entail bringing some of the content currently in the Appendix into the main paper (e.g. Figure 4 plus discussion, and if possible, Figure 11 plus discussion). It will also involve improving the quality of the discussion (emphasising important points), and bringing missing detail regarding the experimental setup (currently in Table 2) to enable more tangible understanding and discussion of the most relevant points

Computational costs: The paper offers limited discussion of the computational cost and scalability of the method, particularly for the Bayesian optimization variant involving Gaussian process surrogates. A more detailed analysis of runtime, memory requirements, and scaling behavior—especially in high-dimensional model spaces—would enhance the practical relevance of the work.

Although we do not provide a comprehensive study of the computational costs and scaling in all settings (model space cardinality, maximum dimension of the model parameters), we do provide a limited study of performance when scaling the model space and parameter dimension simultaneously in Figure 4, Appendix E.1 (as discussed in the response to reviewer ctpC). Since our work builds on IAFs it is more natural to refer to e.g. (Papamakarios et al. 2021) which provides a general overview of the computational costs. As a rule-of-thumb, we found in practice that the size of and choice of transform for the flow composition in $q_\phi$ is subject to the problem type, parameter dimension, and model space cardinality. Problems requiring a highly expressive parameter posterior should use rational quadratic spline transforms in the flow composition. Problems with high parameter dimension may require additional transforms in the composition. Problems with high model space cardinality may require a higher number of blocks in MADE (hidden layers) and/or additional transforms in the flow composition, as well as a context encoder. A brief description of the choice of flow architecture is in Appendix A.2.

For the surrogate-based estimator, we apply incremental, fixed-cost updates over the model space at each time step of the stochastic gradient descent loop. Although a naïve implementation of GPs would incur a cost $\mathcal{O}(t^3 B^3)$ per SGD step $t$, equations 25 and 26 in the appendix describe the updates for the posterior mean and variance of a GP surrogate which we used in practice. In them, the kernel matrix is only evaluated over a $B$-sized batch, hence leading to a $\mathcal{O}(B^3)$ matrix inversion cost per SGD step. Evaluating $M^2$ entries of the posterior predictive variance over a model space with cardinality $M$ costs $\mathcal{O}(B^3 + M B^2 + M^2 B) = \mathcal{O}(M^2 B)$ per SGD step, as $B \ll M$. Over a total of $T$ steps of SGD, we have a cost of $\mathcal{O}(T M^2 B)$ due to surrogate updates. Hence, for moderately sized model spaces, these fixed-cost updates incur a lower cost if we can update the posterior covariance matrix over the entire model space.

We agree these ideas need greater prominence, and so will provide both enhanced and additional discussion as part of moving Figure 4 content (plus discussion) from the Appendix to the main paper.

Mode collapse: Variational inference methods are known to be prone to mode collapse... . It remains unclear whether the proposed approach is susceptible to this issue, and if so, whether any design choices (e.g., CoSMIC flows or masking strategies) help to mitigate it. A discussion or empirical check on this point would be valuable.

This is an important question that we did not address in detail. In extending VI to the transdimensional setting we inherit the same strengths and weaknesses of classical VI, and we make no claim to have eliminated the challenges of mode-collapse. Users of this method would need to take the same steps to manage it in the flows as in the non-transdimensional setting. In the case of mode collapse for the model distribution , we address this (within the limits of the specification of $q_\psi$) with our exploration versus exploitation strategies discussed in Section 3.

Thank you for raising this point. We will include a detailed discussion around this issue in the revised manuscript.

Other applications: Could you provide additional real-world application domains where such distributions naturally arise, beyond the two studied examples?

Trans-dimensional posteriors commonly arise in many other areas, including phylogenetic tree topology search (Everitt et al. 2020), geoscientific inversion (Sambridge et al. 2013), gravitational wave detection (Littenberg and Cornish 2009). The revision will include these citations.

Amortization: How is the variational distribution amortized over the model space $\mathcal{M}$?

As discussed in Section 3, we train one shared recognition network $f_{\phi}$ that (i) is conditioned on model index $m$ to output $q_{\phi}(\theta_{m},u_{\setminus m}\mid m)$ in a single forward pass, and (ii) undergoes heavy optimization during training so that subsequent inference—across any model or dataset—requires only a relatively cheap feed-forward evaluation (no per-model VI runs). We will improve the text in Section 3 to make this clearer.

Conditioning mechanism detail: how is the conditional dependence on the model $m$ incorporated into the formulation of $q_\phi(\theta_m,u_{\setminus m}|m)$?

An equation describing the conditioning mechanism is provided in Appendix A.3, lines 564-571, given inputs $x$ and context $z$, where the context is what we use as the conditioner. This approach is inherited from the nflows package and based on the work in Papamakarios et al. (2017) and Durkan et al. (2019). Additionally, we use a context encoder to enhance expressivity over the model space, described briefly in Appendix A.2 lines 550-552. We appreciate that this could be better explained, and so we will expand the description in the Appendix to more clearly describe the architecture of the context encoder for each experiment.

Tuning complexity: The method relies on careful design of conditional flows (CoSMIC flows) and masking strategies, which may require nontrivial tuning and domain knowledge to generalize to new tasks.

Thank you for raising this point. We respectfully disagree that the design demands non-trivial masking strategies; in our opinion it is substantially less specialized than competing methods. We provide a cookbook approach to designing the masking (Figure 1 and Appendix A.2) enhancing automatability. Extending to other problem types is more straightforward than, for example, designing new RJMCMC proposals. The need for domain knowledge extends only to that required for the set up and specification of a typical Bayesian transdimensional inference problem in the choice of priors and likelihood.

In order to make this clearer for the reader, we will improve the Discussion on this point in the revised version.

References

Everitt, R.G., Culliford, R., Medina-Aguayo, F., Wilson, D.J. 2020. “Sequential Monte Carlo with Transformations.” Statistics and Computing 30: 663–76.

Littenberg, T.B. and Cornish, N.J., 2009. Bayesian approach to the detection problem in gravitational wave astronomy. Physical Review D—Particles, Fields, Gravitation, and Cosmology, 80(6), p.063007.

Sambridge, M., Bodin, T., Gallagher, K., Tkalcic, H. 2013. “Transdimensional Inference in the Geosciences.” Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences 371 (1984): 20110547.

Thanks for your detailed response. The only remaining issue is about the representation of different models $m$, or in your language, the context, which seems important for the conditional dependence on $m$ in the amortization approach. Are the models (or model index) treated as integers here or something else? For examples, when $m$ corresponds to different graphs in DAG learning, how is $m$ represented?

Another question is about the index for the model specific parameters. In the current formulation, it seems that all the models use the same order (i.e., index) for the parameters. As flow-based models (e.g., normalizing flows) are sensitive to the order of parameters, I am not sure if this would lead to extra constraints on the models that can be amortized over. Are the parameters across different models required to have similar alignment before hand?

Are the models (or model index) treated as integers here or something else? For examples, when $m$ corresponds to different graphs in DAG learning, how is $m$ represented?

There is no hard restriction on the representation of $m\in\mathcal{M}$, so long as each model has a unique representation. Proposition 3.1 on line 144 introduces the idea that any discrete distribution $q(m)$ over a finite support has a unique representation as a categorical distribution, i.e. any discrete $m$ where $\mathcal{M}$ is finite can be mapped to a categorical variable. By way of example, in the first example on robust variable selection, $m$ is represented as a binary string that is directly equivalent to $\boldsymbol{\gamma}$, i.e. the included variables are 1, the excluded are 0 in the string. In the non-linear DAG example, the representation of a model is described in Appendix F.7 on line 884 (we link to this appendix on line 206, although we agree the mention of the construction of the model indicator could be made clearer to the reader. We will rectify this in the camera ready version.) In this example, $m=(P,U)$ is a concatenation of the Lehmer code for the permutation matrix $P$ and the flattened binary string from the upper triangular matrix $U$. Additionally, it is permissible to apply a domain-specific transform $m$ before it is input as a context to the flow composition $T_\phi$. In this case, we compute the adjacency matrix $A$ directly from $P$ and $U$ and use that as the context. It should be noted that $m$ is not the context, but rather the model indicator. The context is any function of $m$ that preserves the unique mapping to each model. In the case of DAGs, an adjacency matrix uniquely represents the DAG and is therefore appropriate as the context input to the flow composition $T_\phi$.

We believe this level of flexibility will enable users to map the proposed method to a large variety of inference problems.

Another question is about the index for the model specific parameters. In the current formulation, it seems that all the models use the same order (i.e., index) for the parameters. As flow-based models (e.g., normalizing flows) are sensitive to the order of parameters, I am not sure if this would lead to extra constraints on the models that can be amortized over. Are the parameters across different models required to have similar alignment before hand?

We follow the typical approach in Papamakarious et al. (2021) where flows are composed with reverse permutations between each transformation. This spreads the sensitivity of parameter order approximately evenly and yields good results. We describe the model-specific reverse permutation in Appendix A.2 on line 553. Unfortunately not all of these details could fit in the main paper body but in the camera ready version we will endeavour to link the concepts you have raised to the relevant appendices.