The Free Transformer

## Abstract

Abstract We propose an extension of the decoder Transformer that conditions its
generative process on random latent variables which are learned
without supervision thanks to a variational procedure. Experimental
evaluations show that allowing such a conditioning translates into
substantial improvements on downstream tasks.

## 1 Introduction

Since their invention, the Transformer (Vaswani et al., 2017), and
more specifically the decoder-only Transformers used originally for
the GPT series of models (Radford et al., 2018), have become the core
components of AI systems.

It is remarkable that, after almost a decade, and in spite of
improvements on many aspects of this class of methods, the
autoregressive modelling of Transformers remains essentially
unchallenged. We propose in this paper to revisit this key design
aspect by allowing richer and more natural density models to emerge:

- •
We extend the auto-regressive model of the decoder Transformer
by allowing the conditioning on latent variables, thanks to a
formulation as a conditional Variational Autoencoder
(§ [3.1](https://arxiv.org/html/2510.17558v1#S3.SS1)).
- •
We propose an implementation that requires a very modest
computational and memory usage overhead
(§ [3.2](https://arxiv.org/html/2510.17558v1#S3.SS2)).

The benefits of the proposed method are shown by training 1.5B and 8B
models from scratch and assessing performance on multiple downstream
benchmarks (§ [4](https://arxiv.org/html/2510.17558v1#S4)).

## 2 Motivation

Decoder Transformers are auto-regressive discrete density
approximators. They model a sequence of tokens $S_{1},\dots,S_{T}$ by
estimating the conditional distribution of each given those preceding
it. Sampling is done by generating one token after another, each time
computing the distribution of the next symbol given those generated so
far.

The only density modelling and sampling that such models implement is
that of the generated tokens. In particular, a decoder Transformer
does not make additional latent decisions about the stream of symbols
to generate. Its only decisions are the choices of the tokens
themselves.

Consider, for instance, that we train such a model to generate movie
reviews and that we want to have two clearly separated categories of
negative and positive reviews. Given a large enough model and the
necessary amount of training data, there is no doubt that a decoder
Transformer trained on a dataset of that form would work perfectly and
would generate these two types of reviews. However, to do so, it
would generate tokens one after another and decide, based on the words
generated so far, whether the review it is currently generating is a
positive or a negative one, and continue the process accordingly. In
particular, *the model would not make the explicit decision to
generate a negative or a positive review*. It would produce tokens,
and this notion of a negative or positive review would be implicit in
their posterior probabilities.

Due to the chain rule, any density can be modelled as
autoregressive. However, in particular when the “natural” structure
involves conditioning on latent variables, the autoregressive model
of the signal may be a great deal more complex than the full joint
model including the latent.

We can consider a simple example illustrating that point. Let $Z\sim\mathcal{B}(0.5)$ be a latent “coin flip”, and $X_{1},\dots,X_{T}$ be
equal to $Z$ with independent flips of probability $\epsilon$.

The $X_{t}$s are conditionally independent given $Z$, and we have

$$ $P(X_{t+1}=1\mid Z=z)=\epsilon z+(1-\epsilon)(1-z)$ (1) $$

however, expressed as an auto-regressive model without $Z$, we get:

$$ $P(X_{t+1}=1\mid X_{1}=x_{1},\dots,X_{t}=x_{t})=\frac{\left(\frac{\epsilon}{1-\epsilon}\right)^{{\sum_{s=1}^{t}x_{s}}}(1-\epsilon)^{t+1}+\left(\frac{1-\epsilon}{\epsilon}\right)^{{\sum_{s=1}^{t}x_{s}}}\epsilon^{t+1}}{\left(\frac{\epsilon}{1-\epsilon}\right)^{{\sum_{s=1}^{t}x_{s}}}(1-\epsilon)^{t}+\left(\frac{1-\epsilon}{\epsilon}\right)^{{\sum_{s=1}^{t}x_{s}}}\epsilon^{t}}.$ (2) $$

We could easily come with worse examples when expressed
autoregressively, for instance when the latent variables are positions
in the sequence, e.g. the index where a certain pattern occurs as in
the example of § [4.1](https://arxiv.org/html/2510.17558v1#S4.SS1). What we observe in such
cases is that it requires running estimates of probabilities
(“probability that the target appears here”) for which estimation
errors are unavoidable and problematic.

The consequence is that a purely auto-regressive density model suffers
potentially from several drawbacks:

- •
It requires an unnecessarily complicated computation, and
greater capacity, to implicitly make post-hoc decisions or infer
latent quantities from the generated tokens.
- •
It may be sent off track during the process if, by mistake, a
few tokens generated are erroneous, ambiguous or contradictory with
those generated previously.
- •
Key concepts do not appear spontaneously due to the "natural"
factorization of the distribution, but are built post-hoc by
necessity to fit the training samples better. This may be a
fundamental weakness when operating out of distribution.

The main objective of the present work is to address these issues by
providing the model with the freedom of conditioning its
auto-regressive process on latent random quantities that are not
imposed by the training examples.

For instance, for the review generator example above, the model could
use a random Boolean value to decide once for all whether the tokens
it produces are from the distribution of negative or positive reviews,
removing the need for a complicated posterior estimate from the tokens
already generated.

## 3 Method

Any latent random value $Y_{r}$, whatever its statistical dependency
with the tokens $S_{1},\dots,S_{t}$ and other latent $Y_{1},\dots,Y_{r-1}$ sampled so far, can be expressed under reasonable assumptions
as $f_{r}(S_{1},\dots,S_{t},Y_{1},\dots,Y_{r-1},Z_{r})$ where $Z_{r}$ is a
value coming from a random generator.

Hence, if we provide the model with enough random values $Z_{1},Z_{2},\dots$ sampled independently during generation, a proper training
procedure could in principle build families of latent variables with
arbitrary dependency structure, as long as the model’s capacity allows
it to encode $f_{r}$.

In the same way that the choice of a token during sampling can be
expressed as a function of a random value and the logits, any
activation which is a function of a random value and other activations
can be interpreted as a decision made by the model during the
generative process. Such decisions make the latent activation
non-deterministic functions of the tokens, and observing the latter
only gives a partial information about the former.

Figure: Figure 1: A standard decoder Transformer (a) can be extended to utilize a random state $Z$ in inference (b, left), in which case it has to be trained as a conditional VAE with an encoder (b, right). The Free Transformer (c) reduces the overhead of the encoder by having the random state $Z$ injected in its middle layer (c, left), and using for encoder during training the combination of its first half combined with one non-causal layer specific to the encoder (c, right). See Figure [2](https://arxiv.org/html/2510.17558v1#S3.F2) for a detailed depiction of that architecture.

Figure: Figure 2: The Free Transformer. We omit the normalization layers and residual connections from the model and the batch size from the tensor shapes for clarity. The operators in orange are specific to the encoder and are evaluated only for training or KV cache pre-filling, those with a dashed contour have no trainable parameters. The Binary Mapper is described in § [3.4](https://arxiv.org/html/2510.17558v1#S3.SS4). During generation, the encoder is not evaluated and $Z$ is sampled uniformly among the one-hot vectors of dimension $2^{H}$.

### 3.1 Conditional Variational Autoencoder

Generating a full sequence from scratch with a model that depends on a
random variable $Z$ is trivial: sample $Z\sim P(Z)$ and then run the
standard auto-regressive process, with the computation of the logits
modulated by $Z$.

Training the model, however, is far more involved. Given a training
sample $S$, the objective is to maximize

$$ $P(S)=\int_{z}P(S\mid Z=z)P(Z=z)dz,$ (3) $$

which can be estimated only if we can get $Z$s consistent with $S$,
which amounts to a complex inference problem if we want $Z$ to capture
meaningful structural properties of the sequence.

Providing those $Z$s is the role of the encoder of a Variational
Autoencoder (Kingma and Welling, 2013), whose main purpose is to sample from a
“good” distribution $Q(Z\mid S)$ so that a sampled $Z$ modulates
the decoder in a way that leads it to generate $S$.

We follow this approach and optimize jointly the parameters of the
decoder and the parameters of a second model, which is an encoder in
the VAE sense.

Even though the noise $Z$ has no relation to $S$ initially, if the
training succeeds, the model will use it to structure the generative
process. In the example of a movie review generator of the previous
section, for instance, given a review from the training set, the
encoder would implicitly classify it as positive or negative, and
generate a consistent $Z$. Increasing $P(S\mid Z)$ with that $Z$
could be interpreted as improving the “negative review generator” or
the “positive review generator” that are implicitly encoded in the
decoder’s weights.

A key element of this approach is to limit the amount of information
flowing from the encoder to the decoder through $Z$, so that the
encoder does not provide quantities that should be computed by the
decoder. At the limit the encoder could copy entirely $S$ into $Z$ so
that a trivial decoder, useless without the encoder, hence in
inference, would score perfectly in training.

The formal derivation of the VAE shows that the proper measure of
information is the Kullback-Leibler divergence between $Q(Z\mid S)$
and $P(Z)$, and that the loss to minimize should sum it with the
reconstruction loss, which here is the usual cross-entropy.

### 3.2 Model structure

In what follows, we call “Transformer Block” the usual combination of
a Multi-Head Attention layer and a MLP-like tokenwise module, with
normalisation layers and residual connections.

As pictured on Figure [1](https://arxiv.org/html/2510.17558v1#S3.F1) and [2](https://arxiv.org/html/2510.17558v1#S3.F2),
the Free Transformer is a standard decoder with a noise $Z$ injected
in its middle layer. This allows to share half of the Transformer
block with the encoder, cutting down drastically the computational
overhead by having a single Transformer block that has to be computed
specifically for the encoder. Hence, as we will see, this model
possesses all the components of a decoder Transformer and has an
additional non-causal block and two linear layers for the
encoder. While we did not investigate what is the best depth to inject
$Z$, doing it too early would reduce the encoder’s capacity, and doing
it too late would reduce the decoder’s capacity to process the latent
variables.

For clarity, we omit in what follows the batch size in the tensor
shapes.

As a standard decoder Transformer, the Free Transformer processes a
sequence of tokens by first encoding them with the embedding table
into a tensor $X_{0}$ of shape $T\times D$.

Then it evaluates sequentially the first $L/2$ Transformer blocks to
get $X_{L/2}$ of same shape, and at this point, it samples a sequence
of one-hot vectors $Z=(Z_{1},\dots,Z_{t})\in\{0,1\}^{T\times C}$. During generation, this is done by sampling, for each $Z_{t}$, an
index $c$ uniformly in $\{0,\dots,C-1\}$, and then encoding it as a
one-hot vector of dimension $C$. During training or KV cache
pre-filling, $Z$ has to be consistent with the tokens of $S$ already
fixed, and the sampling is done with the encoder instead, as described
in § [3.3](https://arxiv.org/html/2510.17558v1#S3.SS3).

This tensor $Z$ is processed by a linear layer to obtain a tensor $R$
of shape $T\times D$. Then, the $L/2+1$th Transformer block gets as
input for queries the tensor $X_{L/2}$ and as input for keys and
values the tensor $X_{L/2}+R$. The rest of the Transformer blocks are
evaluated in sequence to get $X_{L}$ which is processed by the read-out
linear layer to obtain the logit tensor $L$ of shape $T\times V$,
where $V$ is the vocabulary suze.

Figure: Algorithm 1 Forward pass of a standard decoder Transformer

Figure: Algorithm 2 Forward pass of a Free Transformer

### 3.3 Encoder and Loss

As stated in the previous section, during training or KV cache
pre-filling, the tensor $Z$ is sampled with the encoder.

The Free Transformer possesses one Transformer block specific to the
encoder, which is non-causal, making the encoder as a whole
non-causal. This is necessary since the conditioning by the decoder
may have long-range effects, requiring the full sequence to be taken
into account to get a proper conditional distribution of the latent.

This encoder-specific block gets as input for the queries a trained
token embedding $\zeta$ replicated to match the sequence length, and
for the keys and values the output of the first half of the decoder’s
blocks. The motivation for using a learned constant input for the
queries instead of the standard representation of the input sequence
is to prevent the encoder from building a token-wise mapping and make it
instead capture global properties of the sequence that may be more
transferable across tasks and data-sets.

A linear readout computes from the encoder block’s output a vector of
dimension $H=16$ for every token. These components are interpreted as
logits of individual bit, used to sample a value in $\{0,\dots,2^{H}-1\}$ which is encoded into a one-hot vector of dimension $2^{H}=65,536$, with gradient pass-through, as described in
§ [3.4](https://arxiv.org/html/2510.17558v1#S3.SS4).

Hence, the random embedding $Z$ is a sequence of $T$ one-hot vectors
$Z_{t}$ of dimension $2^{H}$. The prior distribution used for generation
is uniform $P(Z_{t}=z)=1/2^{H}$, and $Q(Z\mid S=s)$ is the distribution
corresponding to the sampling with the encoder described above. The KL
divergence is then equal to

$$ $\mathds{D}_{\text{KL}}\Big(Q(Z_{t}\mid S_{1},\dots,S_{T})\,\Big\|\,P(Z_{t})\Big)=H\log 2+\sum_{z=1}^{2^{H}}Q(Z=z\mid S)\log Q(Z=z\mid S).$ (4) $$

We control it by adding it to the loss, and prevent its collapse by
using a token-wise free bits method (Kingma et al., 2016). This means
that we sum the KL divergence of individual $Z_{t}$ that are above a
threshold $\kappa$ and ignore the others.

This leads us to use for training loss the sum of the standard
cross-entropy and the following quantity

$$ $\frac{1}{T}\sum_{t=1}^{T}\max\Big(0,\,\mathds{D}_{\text{KL}}\Big(Q(Z_{t}\mid S_{1},\dots,S_{T})\,\Big\|\,P(Z_{t})\Big)-\kappa\Big),$ (5) $$

where the threshold $\kappa$ is an hyperparameter.

### 3.4 Binary Mapper

The last linear layer of the encoder computes for every index $t$ of the
sequence being processed a vector $L_{t}=(L_{t,1},\dots,L_{t,H})\in\mathbb{R}^{H}$, whose components are interpreted as the logits of
individual bits of a binary encoding.

The Binary Mapper samples those bits $B_{t,1},\dots,B_{t,H}$
independentely with

$$ $P(B_{t,h}=1)=\frac{1}{1+e^{-L_{t,h}}},$ (6) $$

and outputs a one-hot vector $Y_{t}$ of dimension $2^{H}$ corresponding to
the resulting value:

$$ $Y_{t,d}=\left\{\begin{array}[]{cl}1&\text{ \ if \ }d=1+\sum_{h=1}^{H}2^{h-1}B_{h,t}\\ 0&\text{ \ otherwise.}\end{array}\right.$ (7) $$

During training, the computation also propagates the gradient of the
probabilities of the $2^{H}$ values. If ${U}(d)=({U}_{1}(d),\dots,{U}_{H}(d))\in\{0,1\}^{H}$ is the binary encoding of $d$, and we
define $G_{t}$ as

$$ $\displaystyle G_{t,d}$ $\displaystyle=P(B_{t}={U}(d-1))$ $\displaystyle=\exp\left(\sum_{h}\log P(B_{t,h}={U}_{h}(d-1))\right)$ $\displaystyle=\exp\left(\sum_{h}(1-{U}_{h}(d-1))\log\left(1-\frac{1}{1+e^{-L_{t,h}}}\right)+{U}_{h}(d-1)\log\left(\frac{1}{1+e^{-L_{t,h}}}\right)\right),$ $$

then the Binary Mapper outputs

$$ $Y_{t,d}+G_{t,d}-\operatorname{detach}(G_{t,d}),$ (8) $$

where $\forall x,\operatorname{detach}(x)=x$ and
$J_{\operatorname{detach}}(x)=0$.

The motivation for using a binary encoding instead of having the
encoder output $2^{H}$ logits directly is to facilitate the gradient
pass-through thanks to the monotonicity of the sigmoid.

## 4 Experiments

We first test the qualitative behavior of the Free Transformer on a
synthetic task in § [4.1](https://arxiv.org/html/2510.17558v1#S4.SS1), then compare it on
multiple benchmarks to baselines with 1.5B and 8B parameters models
for various KL divergence thresholds in
§ [4.4](https://arxiv.org/html/2510.17558v1#S4.SS4), and finally assess the performance
gain of a 8B parameter model trained on 1T tokens in
§ [4.5](https://arxiv.org/html/2510.17558v1#S4.SS5).

### 4.1 Synthetic Dataset

To confirm that the Free Transformer indeed utilizes $Z$ to condition
its generative process, we designed a synthetic dataset and trained a
small Free Transformer with different free-bits thresholds. Doing so
allows to observe what aspects of the modeling are packed by the encoder
in $Z$.

Each sequence in our synthetic training set is generated as follows:

- •
start with a string of 64 underscores “_”,
- •
pick an upper case letter and a position in the sequence at
random, and replace the underscores there with a “target” made of
the selected letter repeated 8 times,
- •
replace any character with an exclamation mark with probability
$1/16$
- •
concatenate a prompt made of the target’s letter followed by a
“>”.

A few sequences generated with that process are shown in
Figure [4](https://arxiv.org/html/2510.17558v1#S4.F4).

We trained a Free Transformer on this data for four different values
of the free bits threshold $\kappa$, and generated with the same
random prompt three groups of sequences with each model, as pictured
in Figure [4](https://arxiv.org/html/2510.17558v1#S4.F4). For each model, in the blue group, the noise
$Z$ is sampled independently for each sequence, whereas we sampled one
$Z$ only for each of the green groups, used to generate all its
sequences.

For very low values of the KL divergence, the model behaves like a
vanilla model (Figure [4](https://arxiv.org/html/2510.17558v1#S4.F4), middle left), and when the
value increases, the model encodes initially the position of the
target alone in the latent state (Figure [4](https://arxiv.org/html/2510.17558v1#S4.F4), middle right),
then encodes both the target position and the noise
(Figure [4](https://arxiv.org/html/2510.17558v1#S4.F4), bottom left), and finally encodes the full
sequence, resulting in incorrect generation (Figure [4](https://arxiv.org/html/2510.17558v1#S4.F4),
bottom right).

Figure: Figure 3: The synthetic sequences of § [4.1](https://arxiv.org/html/2510.17558v1#S4.SS1) are of fixed length, with a “target” made of a random letter repeated 8 times at a random position, an i.i.d. noise of exclamation marks, and a prompt indicating the target’s letter.

### 4.2 Baseline architectures

For assessing performance on standard benchmarks we used decoder-only
Transformers implemented in the same Meta FAIR Transformer codebase as
the one used by Copet et al. (2025) for the Computational World
Model. Those are well optimized models using the SwiGLU non-linearity
(Shazeer, 2020), pre-normalization with RMSNorm
(Zhang et al., 2019), Rotary Positional Embedding (RoPE,
Su et al. 2021), and Group Query Attention (GQA,
Ainslie et al. 2023). The vocabulary size is $2^{17}\approx 130k$.

We used two sizes of models:

- •
A 1.5B model, with $28$ layers, weight tying between the
embeddings and the logit readout, model dimension $1536$, $12$ query
heads, and $2$ key-value heads. It is trained with 47B tokens, which
requires 32 H100s for $\approx$ 12 hours.
- •
A 8B model with the structure of a Llama-3, which is $32$
layers, model dimension $4096$, $32$ query heads, and $8$ key-value
heads. It is trained with 200B tokens which requires 256 H100s for
$\approx$ 24 hours, or with 1T tokens, which takes 5 days.

We compare those baselines to the equivalent Free Transformers, which
require one additional layer for the encoder during training and KV
cache pre-filling, resulting in a compute and memory overhead of $1/28\approx 3.6\%$ for the 1.5B and $1/32\approx 3.1\%$ for the 8B.

### 4.3 Setup and hyperparameters

We kept our findings as clear as possible by avoiding other sources of
performance improvement:

- •
We stuck to the baseline architecture, optimizer, and learning
rate schedule that were used to train the baselines in FAIR’s
framework, and did not optimize any hyperparameter for our setup.
- •
We avoided any recipes for the VAE components, such as removing
sampling in inference. We followed the formal expressions rigorously.
- •
We fixed $H$ to $16$ so that the dimention of $Z_{t}$ was
comparable to the vocabulary size of $2^{17}$.

We stress that the optimization hyperparameters were highly tuned for
the baselines, and it is probable that a combination of an encoder and
a decoder has specific requirements that would greatly benefit from an
adapted training procedure.

**Table 1: Performance of 1.5B models trained on 47B tokens. The training procedure was tuned for the baseline and kept unchanged, but the Free Transformers require $3.6\%$ more compute and parameters for the encoder. See Figure [5](https://arxiv.org/html/2510.17558v1#A3.F5) in Appendix [C](https://arxiv.org/html/2510.17558v1#A3) for the performance during training.**
| 1.5B models (47B tokens) |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Baseline | Free Transformer |  |  |  |  |  |  |  |
|  | 1/4 bit | 1/2 bit | 1 bit | 2 bits |  |  |  |  |  |
| Generative code/math |  |  |  |  |  |  |  |  |  |
| human_eval_plu (pass@1) | 0.055 | 0.079 | +44.44% | 0.079 | +44.44% | 0.085 | +55.56% | 0.085 | +55.56% |
| mbpp (pass@1) | 0.112 | 0.144 | +28.57% | 0.148 | +32.14% | 0.152 | +35.71% | 0.122 | +8.93% |
| gsm8k (em) | 0.025 | 0.028 | +12.12% | 0.027 | +6.06% | 0.033 | +30.30% | 0.027 | +6.06% |
| Multi-choice general knowledge / common sense |  |  |  |  |  |  |  |  |  |
| mmlu (macro_avg/acc_char) | 0.252 | 0.265 | +5.31% | 0.261 | +3.76% | 0.254 | +1.07% | 0.257 | +2.19% |
| csqa (acc_char) | 0.199 | 0.175 | -11.93% | 0.199 | +0.00% | 0.187 | -6.17% | 0.197 | -0.82% |
| hellaswag (acc_char) | 0.593 | 0.591 | -0.40% | 0.594 | +0.15% | 0.592 | -0.27% | 0.595 | +0.32% |
| winogrande (acc_char) | 0.603 | 0.604 | +0.13% | 0.598 | -0.79% | 0.600 | -0.52% | 0.597 | -1.05% |
| obqa (acc_completion) | 0.446 | 0.450 | +0.90% | 0.468 | +4.93% | 0.460 | +3.14% | 0.490 | +9.87% |
| arc_challenge (acc_completion) | 0.400 | 0.392 | -1.93% | 0.386 | -3.43% | 0.405 | +1.29% | 0.385 | -3.65% |
| arc_easy (acc_completion) | 0.596 | 0.602 | +0.92% | 0.592 | -0.64% | 0.603 | +1.06% | 0.592 | -0.71% |
| piqa (acc_char) | 0.734 | 0.736 | +0.22% | 0.738 | +0.52% | 0.734 | +0.07% | 0.733 | -0.15% |
| Multi-choice text understanding |  |  |  |  |  |  |  |  |  |
| race.high (acc_char) | 0.390 | 0.382 | -2.20% | 0.390 | +0.00% | 0.387 | -0.81% | 0.386 | -1.03% |
| race.middle (acc_char) | 0.532 | 0.511 | -3.93% | 0.519 | -2.49% | 0.522 | -1.83% | 0.514 | -3.40% |
| boolq (acc_completion) | 0.583 | 0.632 | +8.39% | 0.614 | +5.35% | 0.648 | +11.12% | 0.620 | +6.29% |
| Culture |  |  |  |  |  |  |  |  |  |
| nq (em) | 0.081 | 0.069 | -15.36% | 0.073 | -9.56% | 0.075 | -7.17% | 0.071 | -11.95% |
| tqa (em) | 0.205 | 0.191 | -6.93% | 0.190 | -7.58% | 0.200 | -2.84% | 0.197 |  |

### 4.4 Exploratory Results

We ran a series of experiments to assess the general behavior of the
Free Transformer, and to calibrate the $\kappa$ threshold.

For any value of $\kappa$, the cross-entropy goes down regularly
during training, with no more instability and spikes than what happens
with the baselines. The KL divergence rapidly goes under $\kappa$ and
stays there. When we compare the cross-entropies for various $\kappa$,
they go down when $\kappa$ increases as expected, but the values
remain extremely close, with a difference of the order of $0.01$ for a
cross-entropy of $\approx 2$ for the 1.5B and $\approx 1.8$ for the
8B.

For both sizes of models, setting $\kappa=4\log 2$, corresponding
to $4$ bits of information per token, resulted in a collapse of the
cross-entropy, indicating that the encoder found a way to channel
fully the tokens to predict, and resulting in a collapse of
performance on the downstream tasks. It is noteworthy that the
baseline 8B model reaches during training a cross-entropy of $1.8=2.59\log(2)$, hence may explain why allowing $2$ bits does not
collapse, while allowing $4$ bits does.

The performance on downstream tasks are given in Table
[4.3](https://arxiv.org/html/2510.17558v1#S4.SS3) for the 1.5B models, and Table
[4.3](https://arxiv.org/html/2510.17558v1#S4.SS3) for the 8B models, both for four different
values of $\kappa$ corresponding to $1/2$ to $2$ bits of information
per token. Graphs of performance during training are given in Appendix
[C](https://arxiv.org/html/2510.17558v1#A3) in Figures [5](https://arxiv.org/html/2510.17558v1#A3.F5) and
[6](https://arxiv.org/html/2510.17558v1#A3.F6).

We observe a substantial increase of performance on HumanEval+, MBPP,
and GSM8K which are arguably the benchmarks requiring some form of
reasoning, and there also is a clear improvement for the 8B model with
1/2 bit of KL divergence on MMLU and CSQA, which are multi-choice
questions.

### 4.5 Results with 1T tokens training

To measure improvement in a more realistic setting, closer to models
actually used in real applications, we trained 8B models on 1T tokens,
which improves drastically the performance of both the baseline and
the Free Transformer.

Given the results with 200B tokens, we chose the value $\kappa=\log(2)/2$ corresponding to half a bit of information per token at
most.

The performance on downstream tasks are given in Table
[4.3](https://arxiv.org/html/2510.17558v1#S4.SS3) and the corresponding graphs during training
in Figure [7](https://arxiv.org/html/2510.17558v1#A3.F7) of Appendix
[C](https://arxiv.org/html/2510.17558v1#A3). We provide in the table the performance
measured at the end of the training as for the other configurations,
but in addition we also give the average over the last third of the
training. We can observe on the graphs that the rate of improvement
tend to be constant on this interval, which justifies averaging to
mitigate the performance fluctuations.

The key result is the boost of performance on HumanEval+, MBPP, GSM8K,
MMLU and CSQA, confirming what we observed in the smaller settings,
and a greater stability on other tasks.

## 5 Previous work

There have been several attempts at combining a VAE and a decoder
Transformer, generally with a focus on improving topic models and
providing ways to guide the generation.

The OPTIMUS model (Li et al., 2020) combines a pre-trained BERT as
text embedding / encoder, with a GPT-2 playing the role of decoder,
which are fine-tuned with a VAE-like loss.

The latent embedding $Z$ is computed thanks to a CLS token, that is by
adding a token to the input and a read-out to extract its embedding in
the output. To modulate the GPT-2 generation with it, it is either (1)
concatenated as an additional token in every layer, or (2) added to
the input token embeddings. Collapse of the KL divergence is prevented
during training with the free bits method (Kingma et al., 2016).

This approach allows for better guided text generation with GPT-2 and
better generalization on low-data languages with BERT.

Xie et al. (2021) extend OPTIMUS with a multi-objective loss, adding in
particular the prediction of the story topic, using the output of
another model as ground truth, to obtain a better embedding space.

The CVAE proposed by Fang et al. (2021) combines two pre-trained
GPT-2, one used as the encoder without causal masking. The embedding $Z$
is an average of the encoder’s output, and the authors propose three
ways to modulate the decoder with linear images of it: (1) add it to
each input token embedding, (2) concatenate it to the Ks and Vs in
every layer, (3) add it before the softmax. Experiments demonstrate
that this method allows controlling the generation without hurting the
quality of the result.

AdaVAE (Tu et al., 2022) is similarly the combination of two
pre-trained GPT-2, the first without causal masking playing the role
of the encoder. The latent embedding $Z$ is extracted from its output
with a slightly modified attention operator. It is then injected into
the decoder by either concatenating an image of it to the keys and
values as in OPTIMUS, or before the softmax as in CVAE.

## 6 Conclusion

The Free Transformer is a direct extension of a standard decoder
Transformer, with the abstract structure of a conditional VAE. It is
implemented with a single additional non-causal Transformer block and
requires a few percent of computational and memory usage overhead.

Its structure makes it able to learn latent random variables
unsupervised, and to condition its generative process on them. In some
ways, this approach aims at achieving in latent space with an
autoencoder what reasoning models do with chains-of-thought in token
space and an RL procedure (DeepSeek-AI et al., 2025). A combination of the
two is, of course, promising.

The performance boost without tuning the optimization hyperparameters
across multiple benchmarks and two sizes of models, is a strong signal
that the overall approach actually improves the inductive bias of the
vanilla Transformer.

Many properties and design choices should be explored. The performance
curves during training are often unstable, possibly due to the
coupling of the optimization of the encoder and the decoder, and using
different optimization methods could be fruitful. The random embedding
itself could take many forms, and the one used in our implementation
is arbitrary.

Finally, the behavior in larger scales, both in parameter count and
dataset size, remains to be investigated.

## References

- Ainslie et al. (2023)
J. Ainslie, J. Lee-Thorp, M. de Jong, et al.
GQA: Training Generalized Multi-Query Transformer Models from
Multi-Head Checkpoints, 2023.
- Austin et al. (2021)
J. Austin, A. Odena, M. Nye, et al.
Program Synthesis with Large Language Models, 2021.
- Bisk et al. (2019)
Y. Bisk, R. Zellers, R. L. Bras, et al.
PIQA: Reasoning about Physical Commonsense in Natural Language, 2019.
- Chen et al. (2021)
M. Chen, J. Tworek, H. Jun, et al.
Evaluating Large Language Models Trained on Code, 2021.
- Clark et al. (2019)
C. Clark, K. Lee, M.-W. Chang, et al.
BoolQ: Exploring the Surprising Difficulty of Natural Yes/No
Questions.
arXiv preprint arXiv:1905.10044, 2019.
- Clark et al. (2018)
P. Clark, I. Cowhey, O. Etzioni, et al.
Think you have solved question answering? Try ARC, the AI2 reasoning
challenge.
In Proceedings of the 2018 Conference of the North American
Chapter of the Association for Computational Linguistics: Human Language
Technologies, Volume 1 (Long Papers), pages 422–435, 2018.
- Cobbe et al. (2021)
K. Cobbe, V. Kosaraju, M. Bavarian, et al.
Training Verifiers to Solve Math Word Problems, 2021.
- Copet et al. (2025)
J. Copet, Q. Carbonneaux, G. Cohen, et al.
CWM: An Open-Weights LLM for Research on Code Generation with World
Models, 2025.
- DeepSeek-AI et al. (2025)
DeepSeek-AI, D. Guo, D. Yang, et al.
DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via
Reinforcement Learning, 2025.
- Fang et al. (2021)
L. Fang, T. Zeng, C. Liu, et al.
Transformer-based Conditional Variational Autoencoder for
Controllable Story Generation, 2021.
- Hendrycks et al. (2021)
D. Hendrycks, C. Burns, S. Kadavath, et al.
Measuring Massive Multitask Language Understanding.
arXiv preprint arXiv:2009.03300, 2021.
- Kingma and Welling (2013)
D. P. Kingma and M. Welling.
Auto-Encoding Variational Bayes, December 2013.
arXiv:1312.6114 [stat].
- Kingma et al. (2016)
D. P. Kingma, T. Salimans, R. Jozefowicz, et al.
Improving Variational Inference with Inverse Autoregressive
Flow, January 2016.
arXiv:1606.04934 [cs].
- Kwiatkowski et al. (2019)
T. Kwiatkowski, J. Palomaki, O. Redfield, et al.
Natural Questions: A Benchmark for Question Answering Research.
In Transactions of the Association for Computational
Linguistics, volume 7, pages 453–466, 2019.
[10.1162/tacl_a_00276](https:/doi.org/10.1162/tacl_a_00276).
- Lai et al. (2017)
G. Lai, Q. Xie, H. Liu, et al.
RACE: Large-scale ReAding Comprehension Dataset From Examinations.
In Proceedings of the 2017 Conference on Empirical Methods in
Natural Language Processing, pages 785–794, 2017.
[10.18653/v1/D17-1082](https:/doi.org/10.18653/v1/D17-1082).
- Li et al. (2020)
C. Li, X. Gao, Y. Li, et al.
Optimus: Organizing Sentences via Pre-trained Modeling of a Latent
Space, 2020.
- Liu et al. (2023)
J. Liu, C. S. Xia, Y. Wang, and L. Zhang.
Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation
of Large Language Models for Code Generation, 2023.
- Mihaylov et al. (2018)
T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal.
Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book
Question Answering.
In Proceedings of the 2018 Conference on Empirical Methods in
Natural Language Processing, pages 2381–2391, 2018.
[10.18653/v1/D18-1260](https:/doi.org/10.18653/v1/D18-1260).
- Radford et al. (2018)
A. Radford, K. Narasimhan, T. Salimans, and I. Sutskever.
Improving Language Understanding by Generative Pre-Training.
Technical report, OpenAI, 2018.
- Sakaguchi et al. (2019)
K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi.
WinoGrande: An Adversarial Winograd Schema Challenge at Scale, 2019.
- Shazeer (2020)
N. Shazeer.
GLU Variants Improve Transformer.
arXiv preprint arXiv:2002.05202, 2020.
- Su et al. (2021)
J. Su, Y. Lu, S. Pan, et al.
RoFormer: Enhanced Transformer with Rotary Position Embedding, 2021.
- Talmor et al. (2019)
A. Talmor, J. Herzig, N. Lourie, and J. Berant.
CommonsenseQA: A Question Answering Challenge Targeting Commonsense
Knowledge.
In Proceedings of the 2019 Conference of the North American
Chapter of the Association for Computational Linguistics: Human Language
Technologies, Volume 1 (Long and Short Papers), pages 4149–4158, 2019.
[10.18653/v1/N19-1421](https:/doi.org/10.18653/v1/N19-1421).
- Tu et al. (2022)
H. Tu, Z. Yang, J. Yang, and Y. Huang.
AdaVAE: Exploring Adaptive GPT-2s in Variational Auto-Encoders for
Language Modeling, 2022.
- Vaswani et al. (2017)
A. Vaswani, N. Shazeer, N. Parmar, et al.
Attention Is All You Need, August 2017.
arXiv:1706.03762 [cs].
- Xie et al. (2021)
Z. Xie, T. Cohn, and J. H. Lau.
Exploring Story Generation with Multi-task Objectives in Variational
Autoencoders, 2021.
- Zellers et al. (2019)
R. Zellers, Y. Bisk, A. Farhadi, and Y. Choi.
HellaSwag: Can a Machine Really Finish Your Sentence?
arXiv preprint arXiv:1905.07830, 2019.
- Zhang et al. (2019)
J. Zhang, R. Xiong, R. Socher, and C. Wang.
Root Mean Square Layer Normalization.
arXiv preprint arXiv:1910.07467, 2019.

## Appendix A Evaluation Benchmarks

- •
HellaSwag: Multiple choices. Common sense focusing on physically situated scenarios. (Zellers et al., 2019)
- •
WinoGrande: Large-scale adversarial Winograd-style pronoun resolution (fill-in-the-blank) designed to reduce annotation artifacts. (Sakaguchi et al., 2019)
- •
ARC (AI2 Reasoning Challenge): Grade-school science multiple choice. (Clark et al., 2018)
- •
PIQA: Physical commonsense multiple choice about everyday goals and affordances. (Bisk et al., 2019)
- •
OpenBookQA (OBQA): Open-book science QA: combines a provided set of core facts with commonsense/world knowledge to answer questions. (Mihaylov et al., 2018)
- •
RACE: Multiple-choice reading comprehension from Chinese middle-school English exams. (Lai et al., 2017)
- •
MMLU: “Massive Multitask Language Understanding”. Questions spanning STEM, humanities, social sciences, etc. (Hendrycks et al., 2021)
- •
CommonsenseQA (CSQA): Multiple-choice QA requiring commonsense relational knowledge (leveraging ConceptNet relations). (Talmor et al., 2019)
- •
BoolQ: Yes/no questions paired with passages to evaluate reading comprehension and entailment-like inference. (Clark et al., 2019)
- •
GSM8K: Grade-school math word problems requiring multi-step arithmetic reasoning. (Cobbe et al., 2021)
- •
HumanEval+: An augmented version of OpenAI’s HumanEval (Chen et al., 2021) with many more unit tests per problem to reduce test fragility and overfitting in code generation evaluation. (Liu et al., 2023)
- •
MBPP: “Mostly Basic Programming Problems.” Short Python programming tasks solvable by entry-level programmers; includes text spec and example tests. (Austin et al., 2021)
- •
NQ: “Natural Questions.” Real user queries paired with Wikipedia pages. (Kwiatkowski et al., 2019)

## Appendix B Performance measures

- •
For generated answers:
–
pass@1 is the proportion of generated pieces of code that
produce the expected behavior when executed.
–
em (“exact match”) is the proportion of generated endings of a
sequence that perfectly match a reference solution.
- •
For multi-choice based on log probabilities:
–
acc_completion is the proportion of correct responses when the
choice is based on the sum of the log probabilities normalized with
the number of tokens of each possible choices.
–
acc_char is the same as acc_completion but normalizes with the
number of characters.
–
macro_avg/acc_char is the average of acc_char over multiple
sub-categories of questions.

## Appendix C Performance during training

Figure: Figure 5: Experiments with 1.5B models trained on 47B tokens. Comparison on standard benchmarks of the baseline and our models. The optimization hyperparameters were tuned for the baseline and kept unchanged, but the Free Transformers require $3.6\%$ more compute and parameters for the encoder.
Refer to caption: x1.png

Figure: Figure 6: Experiments with 8B models trained on 200B tokens. Comparison on standard benchmarks of the baseline and our models. The training procedure was tuned for the baseline and kept unchanged, but the Free Transformers require $3.1\%$ more compute and parameters for the encoder.
Refer to caption: x2.png

Figure: Figure 7: Experiments with 8B models trained on 1T tokens. Comparison on standard benchmarks of the baseline and our models. The training procedure was tuned for the baseline and kept unchanged, but the Free Transformers require $3.1\%$ more compute and parameters for the encoder.
Refer to caption: x3.png

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 4
presentation: 3
questions: Please explain how Figure 1 is related to the specification in Section 3.  What are your variable naming conventions?

Please add the "batch size" to Figure 2.  The answer cannot be "we omit batch size", since this is an essential novelty of your approach over previous latent-variable LLMs.
soundness: 2
strengths: This paper contributes interesting insights into latent variable models for LLMs and a collection of effective methods for making the latent variables work as random choices about future outputs.  Simply adding randomness to hidden representations would not have the same effect, since during LM training these random choices would be independent of future outputs.

A central feature of this latent variable model is that there is one latent vector per token, which is novel with respect to the previous LLM models reviewed in this paper, but related to the work on nonparametric VAEs (Henderson and Fehr, ICLR 2023).  Combined with their method for fixing the KL divergence per token, this means that the number of bits conveyed by the latent variables grows linearly with the length of the generated text, which is essential for modelling the variability in language.

A large number of evaluations are run on reasonably-sized models trained from scratch.  These results succeed in showing the potential of this novel new direction for LLM architectures.
summary: This paper proposes an LLM with one latent vector per token as sources of variation added to a middle layer of an otherwise-standard decoder-only transformer.  As in a VAE, at training time an encoder finds high probability samples of these variables conditioned on the observed text.  Several interesting architectural proposals improve the training efficiency and effectiveness of this model, including sharing the initial layers of the encoder with the initial layers of the decoder, and fixing the KL divergence rather than learning to minimise it.  Empirical evaluations on a large number of LLM benchmarks show promising results, suggesting that adding variance to internal representations of the transformer has added value over just adding variance to the token prediction step.
weaknesses: The paper seems to be work in progress.  There is very little discussion of the empirical results, and no ablation studies other than a standard LLM baseline.  The description of the model misses some key points (see below).

The experiments don't seem to be testing any hypothesis.  The paper reads like they have a cool idea, so lets see what happens. The conclusion is that something interesting happens, but it is not clear what.  This impression is reinforced by the lack of any ablation studies.  

As made clear in Figure 1, there is one latent vector per token.  But Figure 1 is never referenced in the text, and nothing in the specification of the model mentions this fact.  They even say that the prior is "standard Gaussian noise", but Gaussians are parametric distributions, while here the number of sampled values grows with the number of tokens in the text.  This is at best a misleading specification, and I would say incorrect.  In particular, the variable names violate standard naming conventions for vectors versus matrices and tensors.

---

confidence: 5
contribution: 2
initial_rating: 4
presentation: 2
questions: See above.
soundness: 3
strengths: 1. Motivation: The core idea is compelling: enabling models to use explicit latent "plans" rather than relying on purely "post-hoc" autoregressive token-level decisions.

2. Efficiency: The design is highly practical, incurring a small training overhead and zero inference-time cost.

3. Novelty: The "KL Governor" is a creative and new approach to managing the notoriously unstable VAE training objective.

4. Empirical Signals: The strong performance gains on complex coding and math reasoning tasks are significant and warrant attention.
summary: The paper proposes the "Free Transformer," a VAE-based decoder that conditions its generation on a latent variable $Z$. Its core novelty is the "KL Governor," a training mechanism that forces the $D_{KL}(Q(Z|S) || P(Z))$ to a fixed target value $\kappa$ rather than adding it as a loss penalty. This is intended to prevent posterior collapse and, by discarding the encoder at inference, adds no computational overhead. The authors report strong gains on generative tasks like HumanEval+ and GSM8K.
weaknesses: The paper's central claim, while promising, rests on assumptions that would be significantly strengthened by further validation.

1. The Decoder is Not Constrained, and Its Usage of $Z$ is Unproven.

The paper's core thesis rests on the decoder using the latent $Z$. However, the proposed "KL Governor" only constrains the encoder to produce an informative $Z$; it does not place any direct constraint on the decoder. This leaves a critical question unanswered: it is unclear if the powerful autoregressive decoder truly learns to depend on $Z$ or if it simply relies on its own token-level context, effectively ignoring $Z$. This is a key aspect of the posterior collapse problem this paper aims to solve.

To validate this core claim, a critical ablation is necessary: an inference-time test with the latent variable disconnected (e.g., $Z=0$). If the decoder is dependent on $Z$, performance should collapse. Without this experiment, it is difficult to be certain that the reported gains are not just a side-effect of the VAE structure acting as a complex training-time regularizer, rather than the decoder actively using the latent "plan" at inference.

2. Inconsistent Performance and Unaddressed Regressions.

The performance gains appear to be highly task-specific. The paper highlights wins on generative tasks but does not address the notable performance regressions on several standard benchmarks (e.g., -8.46% on boolq and -5.21% on nq for the 8B model). A discussion of this trade-off—why the model might improve on reasoning while regressing on other tasks—is needed to fully interpret the results.

3. Unexplored Design Choices.

The paper introduces several specific design choices without justification or comparison.

(1) KL Governor vs. Alternatives: The "KL Governor" is novel, but the paper would be stronger if it were benchmarked against standard KL regularization methods like $\beta$-VAE or KL annealing, which also aim to prevent posterior collapse.

(2) Architectural Choice: The decision to inject $Z$ only at the middle layer is presented without an ablation or justification. It is unclear why this specific location was chosen over, for example, the input layer or all layers.

---

confidence: 4
contribution: 2
initial_rating: 2
presentation: 2
questions: 1. The encoder design is very simple, so I'm curious if the adding more layers to the encoder can bring better performance. From test-time scaling perspective, I'm expecting to see which factor of the design is important for reasoning capability. Is it still the model size?  If so, then the scaling behavior is the same as GPT based method. Are there any scaling in the latent space?
2. I'm really appreciate if the author can provides more details and experiments based on what I mentioned in the Weakness part. 
3. In real practice, how does the KL setting determined. I'm trying to understand if this is the bottleneck of this method. 
4. Is line 179 a typo?
soundness: 3
strengths: I really appreciate the author for the simple and smart design, which includes minimal changes to the architecture. And according to the paper, the training and inference time doesn't affect a lot by the latent design. 

The model has a clear and explicit latent variable to control the generation, or probably perform some kind of reasoning in latent space. I think this is a great approach and "reasoning in the latent space" is a promising direction.
summary: The paper proposes a “Free Transformer,” a decoder-only transformer that is trained as a conditional VAE but keeps inference almost identical to a normal decoder. The author used two non-causal encoder layers to produce a tokenwise Gaussian posterior, and forces the KL between this posterior and a fixed standard normal prior to be a pre-chosen value. The experiment part provides 2 size of the models performance on several benchmarks.
weaknesses: 1. Model-related issue
* A big autoregressive decoder induces a complex, non-Gaussian posterior, but the paper uses only 2 non-causal layers to output a Gaussian posterior. This is a very narrow variational family and likely leaves an amortization gap. 
* I understand the author want to solve this posterior collapse issue using fixed KL. However, this approach must spend the same KL budget for all samples, so easy examples get unnecessary noise and hard examples cannot request more bits. This directly limits test-time reasoning on harder instances and is tightly coupled to the train–test latent mismatch, which contradict to the idea of this paper of adapting reasoning in different cases via latent variables. 
* Training always feeds structured, sequence-conditioned z at a fixed KL while the inference always feeds unconditioned N(0,I). In my understanding, in conditional generation, posterior should be used instead of prior. The paper does not show how much performance or diversity lost because of this. I hope the author can explain more on this part. 

2. Experimental concerns
* This part is my another major concern due to limited results and details/design included in the paper. 
* Since the author is studying a generative model, then density estimation and perplexity comparison should be included in the experiments. A basic validation perplexity of baseline GPT and other variational/diffusion based methods should be included. Because the method is based on ELBO, it is easy to see the comparison between different baselines in conditional and unconditional generation tasks. Without this, we cannot tell whether the method actually improves language modeling. 
* For the training data, the experiments specify only the token counts (47B, 200B) and the hardware, but not the underlying data distribution (sources, domain mix, code ratio, math ratio, filtering). Since the reported gains are largest on code, math, and CSQA, this is a confound. The authors should either (i) describe the training corpus in detail, or (ii) run ablations where the baseline and the Free Transformer are both trained on an explicitly defined subset (e.g. pure text, pure code) to show that the latent is the cause of the improvement, not the data mixture. 
* There is no concrete numbers to measure the training/inference efficiency. The paper claims “about 6–7% extra cost,” but it does not show per-token TFLOPs, tokens/sec, or memory usage compared to the baseline. Because the main selling point is “latent for free,” this should be quantified. 
* Also, more experiments like prior sampling and diversity results should also be added. The paper does not show this, so we cannot judge whether the learned latent space produces diverse, coherent generations. 
* More analysis on the behavior of this model should be studied. For example, I'm curious about which layers for adding z plays the most important change to the generation. Also, the author didn't show any results of changing z value could lead to different generation, especially the paper has a discussion of examples in the "motivation" section. These should be easily implemented in the ablation study.

---

confidence: 4
contribution: 3
initial_rating: 2
presentation: 2
questions: The authors pretrained the proposed model, what is the perplexity? 

The reported mixed results seems to imply a zero-sum game (at a cost of >5% extra training compute), how would the authors postulate the generic usage of the proposed method?
soundness: 2
strengths: 1. As argued by the authors, variational training of latent-variable LLMs is an interesting and largely under-explored topic. 
2. The authors' proposal of latent inclusion and using KL governor to replace free bits in traditional VAEs are original. 
3. The improvement over baselines "on the generative tasks that put less emphasis on syntax, and measure performance either through the validity of generated code when it is executed (HumanEval+ and MBPP) or through exact matches of the expected response, which requires a complex reasoning (GSM8K)" is inspiring.
summary: This paper proposes to do variational training for decoder-only Transformers. Specifically, the authors introduce a simple way to add some encoder layers (Fig 2) such that the compute from the decoder can be largely reused. The authors also proposed to replace the KL term in the ELBO with a "KL governor", which assigns smaller variance to the posterior if the posterior mean is close to 0, maintaining the KL term at a constant. The authors trained a 1.5B and a 8B model and compared with decoder-only Transformers in a set of reasoning benchmarks.
weaknesses: 1. This paper seems to be written in a rush, with critical technical details missing. For example, in Fig. 2, the input to the encoder seems to be the KV from the middle layer of decoder, and the Q from "a trained token embedding $\zeta$ replicated to match the sequence length" (L211). However, how would these embeddings constitute "Full Self-Att" with those KVs from input tokens given that they are not even the same set of tokens? Do these $\zeta$ have positional embeddings? Another example is that when introducing the details of the "KL governor", which is the key contribution of this paper, the authors only vaguely mention "There is no analytical form for this σ(∥µ∥2), and we need to propagate the gradient through it. So, given the value of the hyperparameter κ we solve it numerically for values of ∥µ∥2 and fit a degree-eight polynomial to it." (L261-263) 

2. The empirical result is mixed, with minimum ablation missing. In the experiment section, the authors only reported the numbers on some reasoning benchmarks. Across all the free-bits value that the author swept over, none of the them outperforms baseline consistently. In fact, >5% degradation is very common in Table 1. IIUC, the baseline is just the autoregressive Transformer, without any RL-based thinking. In terms of ablation, the authors only briefly mentioned the proposed KL governor is more stable than directly optimizing KL (or its variant with free bits), without showing actual results. What if the whole training is deterministic and you still keep the clamp on $\mu$? 

3. Critical related work missing: Phan et al. are probably the first to discuss the relation between latent reasoning and variational inference, who definitely deserve their credits. Kong et al. also introduced variational training to LLMs, which seems to be the most relevant existing work.

Phan et al., Training chain-of-thought via latent-variable inference, NeurIPS'23

Kong et al., Latent Thought Models with Variational Bayes Inference-Time Computation, ICML'25


Decision: Reject