Thank you for your review!

Line 169 - what are these mild assumptions? Should be stated somewhere (or a reference given)

We assume that the initial state distribution and the transition probabilities are such that, for every transient state, there is a non-zero probability of it being sampled in the sequence. We agree with your suggestion and will clarify this in the next version of the manuscript.

Would have liked to see confidence intervals on all results (including e.g., Figure 3 left and Figure 5 left).

Figure 3 (left) and Figure 5 (left) do actually include confidence intervals (95% confidence intervals based on the standard error of the mean computed over 5 and 10 random seeds, respectively), but they are so small that they are barely visible. As far as we are aware, Table 1 is the only float where we present experimental results without confidence intervals (although the corresponding confidence intervals can be read off of Figure 6 in the appendix)

We acknowledge that it is not at all easy to see that confidence intervals are present in the figures you mentioned, and we will experiment with alternative presentation styles (we also welcome your suggestions in that regard).

How does your work compare to Nested Variational Inference https://arxiv.org/abs/2106.11302?

Thank you for bringing this paper to our attention. There are indeed some apparent similarities: their Eq. 7 resembles our Eq. 5. Both involve minimizing divergences across a sequence of distributions. There are also some differences; most importantly, they are concerned with learning a generative model, whereas we learn a discriminative model. Also, it seems to us that the concept of "next step" captures something quite different in both problems.

We are not well-versed enough in the nested importance sampling literature to be able to relate our work to theirs, beyond noticing that the optimization objective appears superficially similar. It seems to us that the two papers solve fundamentally different problems, but we do not exclude the possibility that there are deeper links between these two lines of research.

What kinds of confidence intervals did you use?

We use 95% confidence intervals based on the standard error of the mean (+- 1.96 * stderr).

Figure 5 left - how come 1 token TC outperforms? How is it different from DCE if there's only 1 token? Shouldn't we expect consistency to make more of a difference with more generated tokens? I'm probably missing something here. Is it that you still generate the whole set of tokens, but only evaluate on a cut off set using labels from the whole set?

Your last point is correct: we evaluate the model’s ability to discriminate between correct and incorrect generations given only the first $t$ tokens, but the labels come from the full generation.

Your questions raise two points.

1. When should we expect TC-$\lambda$ to outperform DCE, with few tokens (i.e., short prefixes), or with a full completion? We expect short prefixes to benefit most from temporal-consistency, because that is when the soft targets in Eqs 5 & 7 are most dissimilar from the hard targets used in DCE. In fact, the loss term associated with full sequences is identical for TC-$\lambda$ and DCE (this can be seen by comparing Eq. 2 and the first term in Eq. 5). We were somewhat (pleasantly!) surprised to observe that TC-$\lambda$ models can also outperform DCE models on full sequences.
2. Setting aside differences between TC-$\lambda$ and DCE, how is it possible to achieve a non-trivial ROC AUC with only a single token? This is because, while the model can only see a single token from the generation, it can see the full problem statement. Even though the models are trained on a balanced dataset (every problem has one correct and one incorrect generation in the training data), the models seem to be able to recognize difficult problems statements at test time (and TC-$\lambda$ more so than DCE).

Thank you for your review!

I wonder whether the analysis like convergence of the approach is re-iteration (or minor adaptations) of known properties of TD learning. If so, citations are appreciated.

Indeed, our analysis builds on analogous properties of TD learning. We agree with your suggestion and plan to relate our results more systematically to the corresponding results for TD-learning. These appear in many forms across multiple references, but their presentation in [1] is perhaps the easiest to relate to our developments.

[1] D. Bertsekas, J. Tsitsiklis. Neuro-Dynamic Programming. Athena Scientific, 1996.

Will using the cross entropy loss perform better than the squared loss for LLM generation verification, for example, on the GSM8k studied?

This is a helpful comment, also raised by reviewer mRRN. Although the tasks we consider are naturally framed as classification problems with a binary or a categorical cross-entropy loss, they can indeed also be cast as regression problems with a squared loss. We have addressed this question in more details under mRRN's review (see "Comparison to least squares methods"), including additional experiments comparing our approach to a standard TD-learning approach. We did not have time to run a comparison Qwen2.5 & GSM8K, but we hope that the experiments on the four text classification datasets of Section 4.1. provide some insight into this question.

Why does the difference between the partial sentence classification accuracy (like 16 tokens) and all-token accuracy in Table 1 vary significantly across the datasets?

The datasets in Table 1 span a wide range of characteristics: some have around 20 classes (ohsumed, newsgroups) while others have only 2-4 classes (imdb, ag-news). The informativeness of the first 4 or 16 tokens also seems to vary significantly across datasets, independently of how hard the full-sequence classification task is. In that sense, we are not particularly surprised to see different patterns of performance across datasets.

Does the dataset affect the choices of hyperparameters like lambda?

Yes, hyperparameters do change somewhat across datasets, including $\lambda$. The optimal values are given in Tables 3 & 4 in the appendix. The quantity $1 / (1 - \lambda)$ is the mean of the geometrically distributed weights in Eq. (7), assuming $T \to \infty$, and it has an intuitive interpretation as the average lookahead. In our experiments, the optimal average lookahead ranges between 5 and 50 tokens, depending on the dataset.

Understanding what characteristics of the dataset correlate with optimal values $\lambda$ is an interesting question for future research.

In Line 331: "scaling beyond than 1.3 B parameters remains an open question", this sentence is a bit confusing to me. Do you expect the proposed approach to face scalability issues under larger models (which I don't think so), or do you just mean you did not perform larger scale experiments due to computational constraints? Please enhance the clarity of the sentence.

We mean the latter (we did not perform larger-scale experiments due to computational constraints). Thank you for pointing this out, we will clarify this sentence in the next version of the manuscript.

Thank you. Please update the writing as indicated in the response. I don't have any more concerns about the paper.

Thank you for your review!

If I understand correctly, the theoretical analysis in Section 3 relies heavily on the assumption of working with tabular data, which does not hold for the experiments presented in Section 4. While one might hope that some desirable properties carry over to the parametric setting, the manuscript does not provide strong reasons or intuitions for why this should be the case

We acknowledge that Section 3.1 does not provide much intuition about what makes our temporal-consistency loss particularly effective on the problem in Figure 1, and how this intuition might carry over in more realistic applications with parametric models. We hope to be able to improve this in the next version of the manuscript, given some additional space.

Informally, temporally-consistent methods (such as TC-$\lambda$) exploit information from the dynamics of the sequence. For example, in text classification, one could argue that there are many different ways of organizing words into sentences that capture a similar “semantic state” that is predictive of the class label. In this case, encouraging predictions to be consistent across successive states improves data-efficiency: the model is now implicitly leveraging information from sequences that might have started very differently but reached the same intermediate semantic state. (This data-pooling phenomenon is underpinning the improved statistical efficiency of TC-$\lambda$ over DCE on the problem of Sec. 3.1)

Equation (5) resembles, to me, the typical form of a cost functional in optimal control problems, where the first term corresponds to a terminal cost and the second term resembles the integral of a running cost. However, in line 122, you remark that this equation can be seen as a form of the Bellman equation, suggesting a connection to value functions. Could you please elaborate on this?

In our setup, $p_\theta(\cdot \mid s_t)$ plays the role of the value model. Consider the left-hand side (LHS) and right-hand side (RHS) of Equation (4). Equation (5) can be understood as penalizing the KL-divergence of the LHS relative to the RHS, under an empirical Markov model $p(s_{t+1} \mid s_t)$ given by the data. (Proposition 2 in Sec. 3 makes this precise.)

It is perhaps helpful to make an analogy to a value model in a Markov reward process with discount factor $\gamma$. In that case, the Bellman equation is

And the analogue to our Equation (5) would penalize the squared Bellman error, i.e., the squared difference between the LHS and RHS as

You mention that the computational overhead introduced by the proposed method is negligible, as training time is nearly indistinguishable from the baseline. However, it would be helpful to quantify this claim more precisely, perhaps by providing a theoretical estimate or reporting empirical runtime differences in a controlled setting.

To answer your question, we have re-analyzed a sample of our training runs’ logs. Empirically, the time per training step of TC-$\lambda$ appears to be within 1% of that of DCE. This is explained by the fact that the only overhead of TC-$\lambda$ over DCE is the computation of the pseudo-targets $\{z_t\}$ in (7). This is several orders of magnitude faster than computing forward and backward passes of an LLM. We thank you for your question and will clarify this in the manuscript.

Thank you for your review!

how is TC-lambda different from simply train a value function using TD learning? One might simply use the original TD learning to train such an incremental sequence classification task

This is a helpful comment, also raised by reviewer cjdV. Although all the tasks we consider are naturally framed as classification problems with a binary or a categorical cross-entropy loss, they can indeed also be cast as regression problems with a squared loss. We address this question below in more details, including additional experiments comparing our approach to a standard TD-learning approach.

Is there attempt to anneal the \lambda throughout the training process?

We have not directly experimented with continuous annealing schedules, but we have noticed that the following two-step process (which can be considered as a basic form of annealing) works well:

1. Initialize the classification head by training a linear probe with $\lambda = 1$.
2. Unfreeze the language model parameters and fine-tune end-to-end with the desired value of $\lambda$.

While not strictly necessary to achieve good performance, we have noticed that this leads to very stable and robust training runs that converge quickly. This is the approach we have adopted for our experiments.

Continuously annealing lambda during step 2 above is an interesting idea for future work.

Comparison to least squares methods

In our paper, we train classifiers with a softmax cross-entropy loss. Alternatively, we could frame the classification task as a regression problem (predicting a final scalar reward, which in our case is either $0$ or $1$; in the multiclass case, the $K$ classes can be modeled independently). This regression-based framing naturally suggests using standard value-estimation algorithms from the RL literature, which use a squared loss.

With this in mind, we ran two additional baselines on the four text classification datasets:

• Direct squared loss: regress the final reward directly. This the squared-loss variant of DCE.
• LSTD-$\lambda$: use temporal-difference learning, specifically the offline version of [1]. This is the squared-loss variant of TC-$\lambda$.

[1]: S. Bradtke, A. Barto. Linear Least-Squares Algorithms for Temporal Difference Learning. Machine Learning, 1996.

We compare the predictive performance of these regression baselines on the four text classification datasets of Section 4.1 below.

Key:

• oh: ohsumed
• ng: newsgroups
• im: imdb
• ag: ag-news

Accuracy (higher is better):

Negative log-likelihood (lower is better):

Take-aways:

• methods based on the softmax cross-entropy loss outperform methods based on squared-loss
  • in terms of accuracy, the difference is relatively small
  • the predictive log-likelihood, however, is very poor for regression models (as can be expected).
• as could be expected, the difference is larger for problems with many classes (ohsumed, newsgroups)
• temporally-consistent methods (TC-$\lambda$, LSTD-$\lambda$) tend outperform direct methods (DCE, direct squared loss), especially for short prefixes.

In a way, TC-$\lambda$ provides the best of both worlds: it leverages the benefits of temporal-difference learning while optimizing a loss function that is arguably better suited for classification.