Latent Thought Models with Variational Bayes Inference-Time Computation

Thank you for your thoughtful review acknowledging the novelty of our work. We appreciate your recognition that our claims are supported by clear and convincing evidence. We will address your concerns as follows.

1. Clarification on the cross-attention mechanism.

We wish to humbly clarify a potential misunderstanding regarding the reviewer's comment that incorporating $z$ makes "mask attention in classical autoregressive models not applicable." Our Latent Thought Models (LTMs) adopt the decoder Transformer in the original transformer paper of (Vaswani et al. 2017.), where each layer $l$ of the decoder transformer consists of (1) self-attention sub-layer with attention of $x_t$ to $x_{\leq t}$ (masked), (2) cross-attention sub-layer with attention of $x_t$ to $z_l$, and (3) feed-forward sub-layer. The cross-attention sub-layer (2) for each layer $l$ is ${\rm softmax}\left(\frac{Q_{h_l}K_{z_l}^\top} {\sqrt{d}}\right)V_{z_l},$ where $h_l$ is the hidden embedding of the input $x$ and $z_l$ consists of the latent thought vectors at layer $l$. Crucially, the latent vectors $z$ are inferred by gradient-based variational inference. This means there is no network that maps from $x$ to $z$.

It might be a slight misunderstanding that LTMs "will require $L$ NFEs to evaluate the probability", if $L$ means the number of layers. Given a latent $z$, evaluating $p(x|z)$ only needs 1 NFE. The reviewer correctly notes additional computation is needed, but this is for approximating the posterior $p(z∣x)$ or evaluating the ELBO. This variational inference requires $T_{fast}$ NFEs, scaling with the inference steps $T_{fast}$, not layers $L$.

2. Memory requirements of additional variational posterior model (possibly 2x)?

We wish to respectfully clarify that the 2x memory requirement regarding our additional variational model may be a slight misunderstanding. The classical variational Bayes of our LTMs does not require a separate "variational encoder model" - it only needs the local parameters $(\mu, \sigma)$. The memory overhead for these parameters is almost negligible, with our LTM-L incurring just 1.5% additional memory cost on H100 compared to same-sized AR models with the same batch size. This is more efficient than the encoder (inference model) used in amortized variational inference (VAEs), which typically requires the same memory as the decoder.

3. Demonstration that model learns informative latent vectors (e.g., posterior distribution analysis)

Following your advice, we conduct in-depth analysis of our latent thought vectors through progressive latent inclusion experiments as probing results (link). Our findings reveal: (1) Clear hierarchical latent representation in LTMs, with higher layers integrating information from lower ones. LTM-M shows a significant accuracy jump at the final layer (55% to 100%), while LTM-L demonstrates more gradual information distribution with a notable improvement at layers 9-10 (65% to 95%). (2) Qualitative text reconstruction reveals progression from basic grammatical elements (22% at layers 1-3) to structural patterns (30% at layers 1-6) to coherent semantic content (65% at layers 1-9) to complete sentence (99% at layers 1-10).

4. Computational cost reporting (training time, memory usage)

We reported the trainFLOPs/tok and the number of training tokens in Table 1 of our submission. Here we provide other costs such as training time and memory usage.

All benchmarks were conducted on a single node with 8 H100 GPUs with batch size 512. In our approach, we trade compute for slow learning of global parameters for fast learning of local parameters (line 242-246), which is essential to LTM's data efficiency.

We will revise our paper based on your valuable feedbacks. Thank you!

Thank you for your constructive feedbacks. Below are our responses.

1. Comparison to amortized VI and clarification on multi-step inference

We wish to humbly clarify a possible misunderstanding in terms of the Variational Inference (VI) in our Latent Thought Models (LTMs). LTMs employ the classical variational Bayes (VB) framework rather than amortized inference (as in variational autoencoders (VAEs)).

As detailed in Sec. 2.2 and Algorithm 1 (line 165-183), for each training sequence $x$, the posterior $p(z|x)$ is approximated by a variational distribution $N(\mu, \sigma^2)$. Note that $(\mu, \sigma^2)$ are local parameters specific to each sequence $x$, and are optimized by maximizing the evidence lower bound (ELBO) for each $x$. To clarify the reviewer's comment about "multi-step inference": it refers to the iterative update of $(\mu, \sigma^2)$, where the number of steps is the number of gradient ascent steps for maximizing the ELBO for each $x$.

In contrast to our VB approach, in amortized inference all the training sequences share an encoder (or inference model) with global parameters $\phi$. The encoder maps each $x$ to the corresponding $(\mu, \sigma^2)$ directly. There is no iterative update in typical amortized encoders.

Our ablation study in Sec.3.4 confirmed that classical VB achieves better ELBO and avoids VAE's potential posterior collapse problem.

2. Implementation and hyperparameters

We described our inference algorithm on page 4, Algorithm 1. We discussed the key hyperparameters (inference steps and latent size) in Sec. 3.2. Other hyperparameters can be found in Appx. 1 & 2. To further address your concern, we have built anonymous links for both pseudocode (link) and executable code (link).

3. About novelty of our work regarding variational inference

To our knowledge, our work on incorporating classical VB in language models is novel, as confirmed by Reviewer FqJd and Reviewer ReTi. We are unaware of prior work using classical VI for fine-tuning LLMs. If we are provided with any related references, we would be happy to cite them in revision. The explicit latent abstraction in our model can be an important direction to explore for language modeling.

4. Empirical backing of cognitive analogy of "Language of Thought"

Thanks for your comment. Our model was conceptually inspired by the "Language of Thought" model in cognitive science. However, there has not yet been a machine learning model for it. Our work may be considered the first step towards such a model. To further address your concern, we conduct in-depth analysis of our latent thought vectors through progressive latent inclusion experiments as probing results (link), where our findings reveal:(1) Clear hierarchical latent representation in LTMs, with higher layers integrating information from lower ones. LTM-M shows a significant accuracy jump at the final layer(55% to 100%), while LTM-L demonstrates more gradual information distribution with a notable improvement at layers 9-10(65% to 95%). (2)Case study reveals that lower layers capture grammatical elements and higher layers encode increasingly complex semantic content.

5. Direct analysis confirming latent vectors capture global information

In lines 149-164, the latent vectors $z$ controls the generation of each token in $p(x_{t}|z, x_{<t})$ via cross attention, and $z$ is inferred from the whole sequence $x$: $p(z|x) \propto p(z) \prod_t p_\beta(x_t|z, x_{<t})$, where $\prod_t$ is over all $t$.

Empirically, we validate this by reconstruction experiments where we first infer $z$ from the sequence $x$ and then generate the sequence using only these inferred $z$ (without conditioning on ground truth context). This yields 100% reconstruction accuracy on the OpenWebText validation set, confirming that $z$ successfully captures the global information necessary for generation.

This finding is further supported by our probing results across different models (link is mentioned by response to point 4).

6. Empirical support for conceptual claims (short windows forcing global information into latent variables)

In our response to your point 5 above, $z$ controls the generation of each $x_t$, and is inferred from the whole sequence $x$. With the short context window, the generation of $x_t$ must rely on $z$ to provide information beyond the short window. Empirically, our inferred $z$ can reconstruct the whole sequence accurately as mentioned above. Our probing results (link is mentioned by reponse to point 4) further confirm that the inferred $z$ captures the global information.

We will improve our paper based on your helpful feedbacks. Given our responses above, we humbly request you to reconsider your rating. Thanks!

Thank you for your insightful comments and for recognizing the novelty and strong performance of our work. We address your concerns point-by-point as follows.

1. Scalability.

Our Latent Thought Models (LTMs) scale along two primary dimensions: model size and the number of inference steps ($T_{fast}$). This parallels scaling model size and scaling the length of CoTs in Auto-Regressive (AR) models. Crucially, as discussed in Sec. 2.5, the FLOPs/token cost from $T_{fast}$​ scales linearly, similar to CoT length cost. Therefore, LTMs pose no fundamental asymptotic barriers to scaling beyond those of current LLMs.

In practice, we do acknowledge the engineering challenge of building robust infrastructure for large-scale gradient-based Variational Inference (for pre-training and deployment) and designate this as future work.

2. Inference expense and latency compare to other LLMs

Our analysis in Sec. 2.5 indicates that the inference expense (FLOPs/token) scales linearly with the number of inference steps, $T_{fast}$​. To offer a concrete comparison: performing 16 inference steps with our LTM-M model has latency comparable to generating ~8.7 CoT tokens in the iso-trainFLOPs/tok AR baseline, GPT2-L.

It is also necessary to highlight that practical inference expenses vary significantly depending on the task. While likelihood estimation (Table 1) processes the full sequence length, conditional generation (Fig 6 & Table 2) benefits from typically shorter prefill contexts. Furthermore, for unconditional generation (Table 3), there is no gradient backpropagation in LTMs at inference time, which may lead to considerable speedups compared to other LLM baselines whose decoder networks are significantly larger.

3. Were all compared models using the same GPT2 tokenizer?

Yes. (line 612 in Appendix)

4. Explanation of "lifted latent space"

Thank you for the feedback. "Lifted" is synonymous to "latent" in our context, which is meant to emphasize abstraction from "ground" tokens. We shall make this clear in revision.

5. Typo on page 2: "is not as good as have separate sets"

Thanks. We will fix it in revision.

6. Discussion of other language models like Mixture of Expert(MoE), State Space Model(Mamba) etc.

We will cite these LLM architectures in revision and discuss building our LTMs based on these architectures in future work. Our cross-attention layers that incorporate $z$ can be naturally inserted into these architectures. Since at each cross attention layer, $z_l$ is compact (consisting of only a small number of vectors), the cross-attention is very light, and this is actually quite consistent with the attention-less state vector in Mamba. We shall explore this direction in future work.

We will revise our paper based on your deeply insightful comments. Thank you for your positive evaluation!