Masked Diffusion Models as Energy Minimization

Thank you for your extremely thorough review and constructive suggestions. Below is our point-by-point response.

Weakness 1: Unnecessarily detailed theory

We appreciate this critical perspective. We would like to clarify that although our optimal schedule induced by Condition 3.4 is rather simple, our theoretical framework serves crucial purposes beyond deriving the optimal schedule.

Specifically, MDM's constrained architecture - where both $p_t$ and $Q_t$ are simultaneously determined by $\alpha_t$ - naturally raises questions about its capacity to achieve optimal transport compared to DFM's decoupled optimization. Through Theorem 3.6, we demonstrate that MDM intrinsically encodes optimal rate matrices within its structural limitations. This theoretical guarantee resolves fundamental doubts about MDM's capabilities while explaining its empirical success, thereby strengthening confidence in the framework's mathematical soundness.

Therefore, we believe that our theoretical analysis establishes foundational insights into MDM's operational mechanisms, being helpful in guiding future works.

Weakness 2 & 3: Shifted complexity & Missing comparison of parameterizations

Thank you again for the critical perspective. We would like to clarify the importance of introducing $\gamma_t$ empirically.

First, The $\alpha_t$-$\gamma_t$ duality provides complementary perspectives for schedule design. Common $\alpha_t$s correspond to complex $\gamma_t$s without closed-form expressions (e.g. $\alpha_t=\sin(\frac{\pi}{2}t)$ corresponds to $\gamma_t$ that do not have a closed form). Conversely, Beta-CDF parameterized $\gamma_t$s generally yield non-trivial $\alpha_t$s (e.g. $\gamma_t=CDF(\mathcal{B}(0.9,0.3))$ correspond to nontrivial $\alpha_t$).

More importantly, our empirically effective Beta-CDF tuning method stems directly from the energy perspective. The discovery that standard schedules (linear/sine-squared) correspond to CDF of Beta distribution (Proposition 3.7) motivated systematic Beta-CDF exploration. This connection remains obscure when operating purely in $\alpha_t$ space.

We appreciate your idea in directly parametrizing $\alpha_t$, and although such parameterization might be able to yield empirical improvements as well, the introduction of $\gamma_t$ is still indispensible due to the reasons above.

We hope that our clarification above has addressed your problem.

Weakness 4: Limited impact scope

As clarified in our response to weakness 1-3, despite our rather simple closed-form expression, our theoretical analysis induces unique and nontrivial insights. Also, our effective schedule tuning method, especially the Beta-CDF parameterizing method cannot be done without our theoretical framework. Moreover, as detailed in the response to your next question, we have proposed a more thorough explanation on how to effectively tune the parameters in real-world applications. Such method are also highly motivated by our energy perspective.

Considering these clarifications, we kindly ask you to reassess the significance of our work.

Weakness 5 & Question 2: Lack of a practical algorithm

We appreciate this insightful question. Our theoretical analysis suggests that different downstream tasks may require generated text to possess specific intrinsic structures, which in turn necessitates the generation process to emphasize particular temporal phases. These temporal preferences are captured through different $\gamma_t$ schedules, which induce corresponding optimal $\alpha_t$ through Condition 3.4.

Therefore, we hypothesize that the schedule preferences are mostly inherent to task nature rather than data specifics. To validate this hypothesis, we conducted the following experiments demonstrating that randomly chosen small subsets (50-150 instances) of test data suffice for reliable schedule selection. Specifically, we compare schedule performance between small test subsets and full evaluations (generation length=256, steps=64):

While HumanEval evaluations exhibit greater variance due to smaller test populations (n=164 total), the relative performance rankings remain consistent across subsets - a critical indicator of our method's robustness. This empirical validation confirms that schedule preferences are very probably induced by intrinsic task attributes rather than specific data instances.

Therefore in practice, we recommend conducting initial grid searches using small random test data subsets (about 50~150 instances) across parameters $a,b \in \\{0.1, 0.2, ..., 1.0\\}$ to find a set of well-performed schedules. After the coarse search, we can perform finer-grained selection on shortlisted candidates using larger data subsets. This approach achieves comprehensive parameter space exploration while maintaining computational feasibility.

We will expand this discussion in Section 4 to better guide practitioners in adapting our methodology. We appreciate this valuable question that helps clarify implementation details.

Question 1:

Although I think the presented benchmarks are relatively solid it would be nice to the see evaluation of the proposed method for the optimization of masking schedules of basic MDLM [2] on unconditional generation problems and measure test ppl (e.g. Figure 1 [2]).

Thank you for this question. Our extended experiments demonstrate the broader applicability of our schedule optimization framework to unconditional generation tasks. We still conducted evaluations using Llada due to its unprecedented 8B scale with generation length fixed at 256 tokens and sampling steps at 64. Perplexity measurements were computed via Quen2-7B ([1*]) on 100 generated samples per schedule configuration:

The (0.3,0.9) parameter configuration exhibits catastrophic performance degradation (PPL >4800), producing semantically incoherent outputs that align with its poor performance across multiple tasks in our main experiments. This consistency underscores our method's robustness in identifying non-viable schedules across different generation paradigms.

We hope that our experiment has addressed your problem and further discussion is welcomed.

[1*] Qwen Team. "Qwen2.5: A Party of Foundation Models". 2024.

Question 3:

In both subfigures at Figure 2 there are quite sharp bound after which energy start to drastically grow (green to purple edge). Why is that?

Our computational analysis reveals that schedules deviating significantly from theoretical optima may induce the energy functionals to dirverge. The Beta distribution parameters systematically govern critical schedule characteristics including concavity-convexity properties and inflection point locations. When these parameters stray beyond critical thresholds from their optimal configurations, the energy functionals will may diverge to infinity.

Our large-scale experiments corroborate this phenomenon: certain schedule may exhibits catastrophic performance degradation on specific tasks while maintaining normal functionality on others. This bifurcation originates from the energy landscape's acute sensitivity to $\gamma_t$ variations. Different $\gamma_t$ schedules create distinct energy topographies where non-optimal configurations reside in high-energy areas, leading to suboptimal generation performance.

Question 4:

In Figure 2 caption the following said “Axes represent the beta-parameterization of $\alpha_t$ (see Sec. 3.3)”. What exactly does it mean, what is beta-parameterization of $\alpha_t$ ? I could not find such information in Sec. 3.3

Thank you for highlighting this presentation gap. The beta-parameterization refers to constructing $α_t$ schedules through the Beta CDF: $α_t = CDF(\mathcal{B}(a,b))(t)$, where $(a,b)$ parameterize the Beta distribution. Figure 2 illustrates how these parameters influence energy landscapes under fixed $\gamma_t$ schedules through Condition 3.4. Each axis coordinate corresponds to specific $(a,b)$ pairs that generate distinct $α_t$ via this CDF mapping. The visualization demonstrates how Condition 3.4-guided parameter selection identifies energy minima. We will clarify this methodology explicitly in revisions to prevent misinterpretation.

Question 5:

In Def 2.1 the unmasking process rate matrix is used, while as I understand it there also could be masking process rate matrix, am I right? If that is the case and in general these energy formulations (Def 2.1, Def 2.2, Def 2.3) can be defined in both ways (forward and backward) I think it should be mentioned somewhere in the main part of the paper.

Thank you for your suggestion. While energy formulations could indeed be constructed for both masking and unmasking processes theoretically, our framework deliberately focuses on the unmasking rate matrix since the unmasking process directly governs the practical sampling trajectory – the primary optimization target for real-world applications.

We hope that our response above has addressed your concerns, and we welcome further discussion if needed.

I thank the authors for answering my questions. They have addressed most of my concerns. I would like to raise my rating from 4 to 5.

Thank you for the recognition of our contributions and the thoughtful comments. Below is our point-by-point response.

Weakness 1:

However, the optimal sampling schedule defined to minimize the the three formulations that the authors proposed sounds not to improve anything because the authors also showed the "Invariance of Training Loss to the Mask Schedule" in Appendix C. So, I wonder if the three are equivalent formation of the same concept.

We appreciate this important question. The three energy formulations proposed in our work quantify the sampling process's "energy cost" from distinct perspectives, which fundamentally differ from the training loss function. To be more specifically, The training phase aims to learn a neural network approximation of the target distribution whereas the sampling phase addresses the operational trajectory for sample generation from the learned distribution.

Therefore, while Appendix C demonstrates training loss invariance to mask schedules – enabling post-training schedule tuning – the sampling process remains highly schedule-dependent, as evidenced by our empirical results.

We will adapt the above clarification in the main text of the final version to make our theoretical framework clearer.

Weakness 2:

Also, in the three definitions, what if the denominator, say in the eqn. (5), equals 0?

We thank you for this meticulous observation. We would like to clarify that such cases are natually excluded from the summation, therefore does not affect our framework. Consider the kinetic energy definition:

$\mathcal{E}_ k(p_t,Q_t;\gamma_t)=\mathbb{E}_ {t,p_t(z)}\sum_ {x:x\neq z}\frac{1}{\dot{\gamma}_tp_t(x)}Q_t(z,x)^2.$

To calculate this formulation, we first sample $z$ that have positive probability, i.e. $p_t(z)>0.$ Due to the detailed balance condition, as long as $Q_t(z,x)>0$, we must have $p_t(x)>0$ due to the continuity of the probability distribution. Therefore, the case $p_t(x)=0$ can only occur when $Q_t(z,x)=0$.

Since the summation is of all probability flows between states, such $(z,x)$ pairs carrying zero transition probability ($Q_t(z,x) = 0$) are inherently and natually excluded from the summation, as it suggests that the flow between $z$ and $x$ does not exist at all. Our theoretical framework and proofs implicitly incorporate this exclusion principle throughout the analysis. We acknowledge the need to explicitly state this detail and we will revise accordingly in our final version in caption to enhance clarity. However, such concern does not affect our theoretical framework. Thank you for highlighting this important presentation consideration.

Weakness 3:

I am further concerned when I read Page6, Figure 2: according to the pictures, the color indicates energy values, it seems like discontinuous for Beta(a,b). This suggested the energy function defined did not consider possible 0 in the denominator.

We appreciate your careful examination. The apparent discontinuity in Figure 2 stems from our discrete parameter sampling strategy for visualization purposes rather than inherent energy function discontinuity. Energy values actually vary continuously with respect to $(a,b)$ parameters. Figure 2 is to emphasize the relationship between different weight functions and corresponding energy minima positions, therefore we only selected sparsely distributed $(a,b)$ pairs to show the coarse relationship. As addressed in our previous response regarding denominator zeros, this potential concern does not affect the theoretical validity or experimental results presented.

Question 1:

It seems there is some mistake when the authors introduced the MDM. Page 2: line 62: The authors mentioned: "a masking process (operating backward in time, t = 1 → 0) and its reverse, the unmasking process (t = 0 → 1) ...", this seems opposite to what was in the previous literature where forward diffusion is masking and backward is unmasking. As a result, in Line 57: The authors assume p0(z) follows a uniform distribution, but they did not specify what is the target distribution p_1(z). Again, this contradicts with existing literatures. For example, ref[8] of the paper discussed the forward diffusion to a uniform distribution; and ref [43] works on the case that x_1 follows a Dirac distribution.

Thank you for this insightful observation. Our temporal parameterization follows conventions from energy-based sampling literature [1*,2*] rather than standard diffusion frameworks. Specifically: $t=0$ denotes the initial sampling state (fully masked Dirac distribution in MDM) and $t=1$ marks the final generated state.

This orientation aligns the forward temporal direction with the unmasking process – the core focus of our energy analysis. The reversed timeline offers two key advantages: 1) It simplifies energy functional expressions by synchronizing time progression with generation advancement, and 2) Ensures mask/interpolation schedules naturally increase from 0 to 1. While differing from some diffusion literature, this convention maintains internal consistency and better serves our analytical framework's objectives.

[1*] Neta Shaul et al. “On Kinetic Optimal Probability Paths for Generative Models”. In: International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA. Ed. by Andreas Krause et al. Vol. 202. Proceedings of Machine Learning Research. PMLR, 2023, pp. 30883–30907.

[2*] Neta Shaul et al. “FlowMatching with General Discrete Paths: A Kinetic-Optimal Perspective”. In: CoRR abs/2412.03487 (2024). DOI: 10.48550/ARXIV.2412.03487. arXiv: 2412.03487.

Weakness 4 & Question 3:

While authors focused on Q_t(z, x), I suggest they should provide p_t(z, x) when defining MDM.

The paper did not explicitly mentioned what the target distribution is; Figure 4 showed some, but did not explicitly explained it.

We appreciate this technical clarification. In continuous-time Markov chains, the marginal distribution p_t(z) evolves via the Kolmogorov forward equation $\frac{dp}{dt} = p_tQ_t$. For MDM's specific $Q_t(z,x)$ structure, $p_t(z)$ generally lacks closed-form solutions but remains fully determined through $Q_t$ specifications. Therefore following MDM traditions [1*,2*], we only specify $Q_t(z,x)$ and $p_t(z)$ is then fully determined.

As for the target distribution, our main experiments focus on practical scenarios where target distributions correspond to task-specific data distributions without explicit analytical forms. As for the synthetic experiment in Figure 4, the target distributions indeed have analytical forms and details are presented in Appendix E.4.

[1*] Jacob Austin et al. “Structured denoising diffusion models in discrete state-spaces”. In: Advances in Neural Information Processing Systems 34 (2021), pp. 17981–17993.

[2*] Jingyang Ou et al. “Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data”. In: ICLR. 2025.

Question 2:

Page2: Eqns. (1) (2) did not make sense with (z not equal to x) in between the two lines, seems like a mistake.

Thank you for this meticulous observation. The $z \neq x$ condition rigorously reflects continuous-time Markov chain conventions [1*]. Off-diagonal elements $Q_t(z,x) (z \neq x)$ & represent transition rates between distinct states, while diagonal elements satisfy $Q_t(z,z) = -\sum_{x\neq z}Q_t(z,x)$ to conserve probability mass. Explicitly specifying $z \neq x$ ensures mathematical precision when defining transition dynamics, as diagonal elements derive from off-diagonal terms rather than being independently specified.

[1*] Jingyang Ou et al. “Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data”. In: ICLR. 2025.

We hope that our response above has addressed your concerns, and we welcome further discussion if needed.

Thank you for recognizing our contributions and providing thoughtful feedback. Please find our response below.

Weakness & Question:

While experiments show improved sample efficiency, the performance gains are limited to a subset of benchmarks (notably, not significant on BBH or GSM8K)

The proposed method demonstrates strong improvements in few-step sampling on HumanEval and Hendrycks Math. However, the gains are less significant or inconsistent on other benchmarks like BBH and GSM8K. Can you elaborate on why the proposed energy-based schedule performs better on some tasks than others?

We appreciate this insightful question. We clarify that our methodology provides both a theoretical foundation for understanding schedule preferences and a practical mechanism for task-specific optimization, rather than proposing a universally superior schedule.

In fact, linear schedule $\alpha_t = t$ is a special case in our framework since it corresponds to $\gamma_t=CDF(\mathcal{B}(0.5,0.5))$, a point in our parameter space. Therefore, the empirical parity on certain benchmarks confirms that linear schedules already serve as near-optimal candidates for specific task categories. However, our framework's value lies in enabling systematic task-specific schedule optimization, overcoming the limitations of previous one-size-fits-all approaches.

Specifically, the linear schedule $\alpha_t = t$ corresponds to $\gamma_t=CDF(\mathcal{B}(0.5,0.5))$, which has higher derivatives on both initial and final generation phases while being relatively static in the middle. Since $\dot{\gamma}_t$ appears in the denominator of our energy functional expression, this implies that tasks requiring sustained refinement throughout generation (particularly middle phases) rather than the starting or ending phases might inherently favor the linear schedule.

Among the six tasks we experimented on, BBH is only one that focuses on general reasoning problems. As for GSM8K, although it is mathematics-focused, its prompt requires the generated text to use detailed natural language explanations compared to other math benchmarks. This might be the reason why sustained refinement in the middle phases are preferred. Considering this similarity, it is reasonable that GSM8K and BBH have similar preference on schedules although it is difficult to mathematically deduce the exact expression of the energy functionals.

Considering the complexity of real-world tasks, it is indeed a challenging direction to systematically study the relationship between tasks and their prefered schedules. A promising research direction involves designing intermediate benchmarks that balance practical relevance with analytical tractability - more sophisticated than our current toy examples yet simpler than real-world tasks.

We will incorporate this analysis in Section 4 of our paper in the final version. We appreciate this insightful question and we hope our response addresses your concerns. Further discussion is welcomed if needed.

Thank you for your careful review and thoughtful comments. Below is our point-by-point response.

Weakness 1:

Although the study introduces a kinetic viewpoint, the proposed masking schedule relies heavily on geodesic insights already explored in [1]. This makes the kinetic formulation feel more like a minor addition, with the main contribution closely following prior work—thereby limiting the novelty of this work.

The unified theoretical framework appears somewhat disconnected from the proposed schedule. The most practical contribution lies in the new masking schedule, which is primarily inspired by prior work—limiting the overall novelty.

We appreciate this critical perspective. We would like to clarify that Our framework fundamentally operates within the optimal transport paradigm for efficient sampling beyond the geodesic perspective. As optimal transport connects directly with efficient sampling [1*,2*], we establish provably optimal schedules for sampling efficiency in masked diffusion models. More importantly, unlike DFM's decoupled optimization of $p_t$ and $Q_t$, MDM's unique architecture – where both $p_t$ and $Q_t$ are jointly determined by $\alpha_t$ – natually raises fundamental questions about whether such a constrained setting still can reach optimality. Our analysis demonstrates that MDM intrinsically encodes optimal rate matrices within its constrained structure. This theoretical guarantee strengthens confidence in MDM's elegant framework, particularly given its empirical success. Such a perspective is absent in [1].

Technically, our core objective of our analysis is to minimize the kinetic and conditional kinetic energy, which cannot be interpreted easily within the geodesic perspective. Therefore, the geodesic method only gives us a mathemetical tool rather than a core goal. By utilizing the geodesic perspective, our main theorems are successfuly proven and understanding of MDM's capabilities are deepened. As noted by Reviewer dGWj, such synthesis "presents a fresh combination of concepts, extending prior research in non-trivial ways."

Even within the geodesic perspective, it is unclear how the pure geometric method in [1] can be used to prove the optimality of Condition 3.4. By introducing the geodesic energy, our proof of Lemma 3.5 is proved using functional minimization methods, not only extending the original proof, but also successfully connected to our main goal in optimal transport.

Therefore, although the closed-form condition was hinted in [1], our theoretical contributions beyond prior work is still non-trivial. We will clarify these points in the revised version and we hope that this response has addressed your question.

[1*] Cédric Villani. Topics in optimal transportation. Vol. 58. American Mathematical Soc., 2021.

[2*] Neta Shaul et al. “On Kinetic Optimal Probability Paths for Generative Models”. In: International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA. Ed. by Andreas Krause et al. Vol. 202. Proceedings of Machine Learning Research. PMLR, 2023, pp. 30883–30907.

Weakness 2:

The use of the Beta CDF as the masking schedule $\gamma$ is reasonable, as it generalises existing schedules such as linear and squared sine. It also enables inference-time optimisation in MDMs. However, the choice of the Beta distribution parameters $a, b$ appears to be heuristic. Could the authors elaborate on how to systematically search for or optimise $a, b$ rather than relying on trial and error?

The choice of parameters $a, b$ in the Beta distribution across different settings lacks sufficient justification.

We appreciate this insightful question. Our theoretical analysis suggests that different downstream tasks may require generated text to possess specific intrinsic structures, which in turn necessitates the generation process to emphasize particular temporal phases. These temporal preferences are captured through different $\gamma_t$ schedules, which induce corresponding optimal $\alpha_t$ through Condition 3.4.

Therefore, we hypothesize that the schedule preferences are mostly inherent to task nature rather than data specifics. To validate this hypothesis, we conducted the following experiments demonstrating that randomly chosen small subsets (50-150 instances) of test data suffice for reliable schedule selection. Specifically, we compare schedule performance between small test subsets and full evaluations (generation length=256, steps=64):

While HumanEval evaluations exhibit greater variance due to smaller test populations (n=164 total), the relative performance rankings remain consistent across subsets - a critical indicator of our method's robustness. This empirical validation confirms that schedule preferences are very probably induced by intrinsic task attributes rather than specific data instances.

Therefore in practice, we recommend conducting initial grid searches using small random test data subsets (about 50~150 instances) across parameters $a,b \in \\{0.1, 0.2, ..., 1.0\\}$ to find a set of well-performed schedules. After the coarse search, we can perform finer-grained selection on shortlisted candidates using larger data subsets. This approach achieves comprehensive parameter space exploration while maintaining computational feasibility.

We will expand this discussion in Section 4 to better guide practitioners in adapting our methodology. We appreciate this valuable question that helps clarify implementation details.

Weakness 3:

Additionally, how can varying $a, b$ values lead to improvements over the linear and squared sine schedules in certain scenarios (especially in low-step sampling regimes)? Any insights or empirical observations would be appreciated.

Thank you for this insightful question. As mentioned in the response to Weakness 2, different intrinsic structures of tasks may induce different weight function $\gamma_t$, therefore inducing different $\alpha_t$ preferences. Since $\dot{\gamma}_t$ appears in the denominator of our energy functional expression, temporal phases where $\gamma_t$ has larger derivatives get less weight and phases where $\gamma_t$ is relatively flat get more focused. By varying the parameter $(a,b)$, we fundamentally alter the schedule's concavity and inflection points, therefore changing the temporal focus. When the temporal focus is aligned with intrinsic demand of the specific task, sampling performance is thereby promoted.

Take the first toy experiment in Figure 4 as an illustrative case. Since the target distribution only have positive probability on extreme points (all zeros or all ones), precise trajectory control becomes paramount in initial generation phases to prevent simultaneous 0-1 generation. As the generation proceed to the ending period, however (for example 4 zeros have already been unmasked), the distribution itself (which is parameterized by a neural network in real world experiments) will prohibit the generation of ones since it will lead to invalid combinations. This eliminates the need for intensive late-stage focus, enabling efficient resource allocation to the decisive initial phases. Such temporal preference make the model prefer the parameter (0.6,0.3), where the starting period of generation is more weighted.

In low-step regimes, strategic allocation of limited computational resources becomes more critical. This explains why our tuning method lead to improvements especially under low step sampling scenarios.

We hope that our response above has addressed your concerns, and we welcome further discussion if needed.

Thank you for your extremely thorough review and acknowledgment of our contributions. Below is our point-by-point response.

Weakness (Quality) & Question 1:

While mathematically sound, the authors don’t explicitly discuss the connection between the assumed continuous-time diffusion setting in the theoretical part and the finite-step discrete sampling in real settings.

The theorems assume a continuous-time interpolation from 0 to 1. In practice, a discrete sequence of masks must be chosen. Is the sine-squared schedule still exactly optimal in a discretized sense, or only in the continuous limit? Currently this is only done implicitly by tuning Beta parameters for a fixed step count. A brief clarification of the theoretical connection would be helpful.

We appreciate this insightful question. Indeed, under Riemann discretization, MDM trajectories still minimizes a corresponding discrete energy functional that converges to the continuous formulation as the number of steps $N \to \infty$.

Due to the equivalence of the three energy functionals, we focus our analysis on the geodesic energy. Recall that $\mathcal{E}_ g$ can be equivalently expressed as:

$\mathcal{E}_ g = \int_{0}^1 \frac{||\dot{y}_t||^2}{\dot{\gamma_t}} dt.$

Therefore, we can discretize $\mathcal{E}_g$ as:

$\mathcal{E}_ g^N = \sum_{n=0}^{N-1} \frac{||y_{n+1}-y_{n}||^2 /(\Delta t)^2}{|\gamma_{n+1}-\gamma_{n}| /\Delta t} \Delta t = \sum_{n=0}^{N-1} \frac{||y_{n+1}-y_{n}||^2 }{|\gamma_{n+1}-\gamma_{n}|},$

where $\Delta t=1/N.$ For simplicity, consider the special case where $\gamma_t = t$. Then we have $|\gamma_{n+1}-\gamma_{n}| = \frac{1}{N}$. We now demonstrate that the $N$ equidistant points along the geodesic still minimize this functional. First, we note that the squared chord length on the unit sphere satisfies $||y_{n+1}-y_{n}||^2 = 2(1-\cos(\theta_n))$, where $\theta_n$ denotes the angle between $y_n$ and $y_{n+1}$. This transforms the minimization problem into finding $N-1$ intermediate points along the great circle arc between $y_0$ and $y_1$ that minimize $\sum_n 2(1-\cos(\theta_n))$.

From the properties of spherical geodesics, we know that placing all intermediate points along the geodesic path enforces the constraint $\sum\theta_n = \frac{\pi}{2}$. Given that $1-\cos(x)$ is convex over $(0, \pi/2)$, Jensen's inequality establishes:

$\sum_n 2(1-\cos(\theta_n)) \geq N\cdot 2(1-\cos(\frac{\pi}{2N})),$

with equality achieved when points are uniformly distributed along the geodesic. Any deviation from the geodesic path would further increase the cumulative angular separation $\sum\theta_n$ beyond $\frac{\pi}{2}$, consequently raising the total energy. This proof extends to arbitrary $\gamma_t$ schedules through proportional point distribution based on $|\gamma_{n+1}-\gamma_{n}|$.

We will add the above discussion in Section 3.2 of the paper in the final version. We are grateful for this constructive question that allows us to strengthen the theoretical grounding of our approach.

Weakness (Clarity):

While the notation is mostly standard, expressions like in the conditional kinetic energy definition (Eq. 6) might require readers that are not familiar with OT to study the appendix thoroughly before being able to follow the results and keeping track of the notation.

In addition, some key ideas, like the mask schedule invariance of the ELBO objective are mentioned somewhat briefly and more emphasis on such points in the main text could improve the reader’s understanding of why post-training schedule tuning is possible.

Thank you for the constructive suggestions. We will improve the presentation by providing clearer explanations of the definitions in the main text to enhance readability. Also, the ELBO's invariance to mask scheduling indeed plays a crucial role in our framework, therefore we will emphasize this fundamental property more prominently in the main text to improve conceptual clarity.

Question 2:

How are the parameters of the Beta-CDF selected for a given task? Did you perform a grid search on (a, b) and pick the best performing schedule? If yes, what are the parameters of this search? It would be useful to know the number of evaluations needed and whether it is feasible to apply to very large models. Discussing the tuning strategy would also help readers to adapt your method more easily and extend this part of it.

We appreciate this insightful question. Our theoretical analysis suggests that different downstream tasks may require generated text to possess specific intrinsic structures, which in turn necessitates the generation process to emphasize particular temporal phases. These temporal preferences are captured through different $\gamma_t$ schedules, which induce corresponding optimal $\alpha_t$ through Condition 3.4.

Therefore, we hypothesize that the schedule preferences are mostly inherent to task nature rather than data specifics. To validate this hypothesis, we conducted the following experiments demonstrating that randomly chosen small subsets (50-150 instances) of test data suffice for reliable schedule selection. Specifically, we compare schedule performance between small test subsets and full evaluations (generation length=256, steps=64):

While HumanEval evaluations exhibit greater variance due to smaller test populations (n=164 total), the relative performance rankings remain consistent across subsets - a critical indicator of our method's robustness. This empirical validation confirms that schedule preferences are very probably induced by intrinsic task attributes rather than specific data instances.

Therefore in practice, we recommend conducting initial grid searches using small random test data subsets (about 50~150 instances) across parameters $a,b \in \\{ 0.1, 0.2, ..., 1.0 \\}$ to find a set of well-performed schedules. After the coarse search, we can perform finer-grained selection on shortlisted candidates using larger data subsets. This approach achieves comprehensive parameter space exploration while maintaining computational feasibility.

We will expand this discussion in Section 4 to better guide practitioners in adapting our methodology. We appreciate this valuable question that helps clarify implementation details.

Weakness (Significance) & Question 3:

Nevertheless, the empirical performance gains are not universal, as some benchmarks showed no improvements over linear schedules. This shows that further research is needed to gain a better understanding, which the authors acknowledge as a future direction.

Do you have hypotheses why GSM8K and BBH did not benefit from optimized Beta schedules? Is it because those tasks require very specific refinement at the end, so a linear schedule is already as good as it gets? Providing some discussion on these cases would be valuable.

Our theoretical analysis suggests that different tasks may prefer different schedules, therefore the observed parity with linear baselines in certain tasks aligns with our theoretical framework.

Specifically, the linear schedule $\alpha_t = t$ corresponds to $\gamma_t=CDF(\mathcal{B}(0.5,0.5))$, which has higher derivatives on both initial and final generation phases while being relatively static in the middle. Since $\dot{\gamma}_t$ appears in the denominator of our energy functional expression, this implies that tasks requiring sustained refinement throughout generation (particularly middle phases) rather than the starting or ending phases might inherently favor the linear schedule.

Among the six tasks we experimented on, BBH is only one that focuses on general reasoning problems. As for GSM8K, although it is mathematics-focused, its prompt requires the generated text to use detailed natural language explanations compared to other math benchmarks. This might be the reason why sustained refinement in the middle phases are preferred. Considering this similarity, it is reasonable that GSM8K and BBH have similar preference on schedules although it is difficult to mathematically deduce the exact expression of the energy functionals.

Considering the complexity of real-world tasks, it is indeed a challenging direction to systematically study the relationship between tasks and their prefered schedules. A promising research direction involves designing intermediate benchmarks that balance practical relevance with analytical tractability - more sophisticated than our current toy examples yet simpler than real-world tasks.

We will incorporate this analysis in Section 4 of our paper in the final version. We appreciate this insightful line of questioning and we hope our response addresses your concerns. Further discussion is welcomed if needed.