GPT-4 Turbo: Research Preview (128K token limit, Short-Term Availability)
If this app doesn't respond, it's likely due to our API key hitting the daily limit of our organization. Consider trying our GPT-3.5 app:
https://huggingface.co/spaces/yuntian-deng/ChatGPT
Chatbot
explain the following like i am a layperson: ““LORD: LOW RANK DECOMPOSITION OF
MONOLINGUAL CODE LLMS FOR ONE-SHOT COMPRESSION
Ayush Kaushal
Universite de Montr ´ eal, Nolano AI ´
<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>
Tejas Vaidhya
Mila, Universite de Montr ´ eal, Nolano AI ´
<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>
Irina Rish
Mila, Universite de Montr ´ eal, Nolano AI ´
<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>
ABSTRACT
Low Rank Decomposition of matrix - splitting a large matrix into a product of
two smaller matrix offers a means for compression that reduces the parameters
of a model without sparsification, and hence delivering more speedup on modern hardware. Moreover, unlike quantization, the compressed linear layers remain
fully differentiable and all the parameters trainable, while being able to leverage
the existing highly efficient kernels over floating point matrices. We study the potential to compress Large Language Models (LLMs) for monolingual Code generation via LOw Rank Decomposition (LoRD) and observe that ranks for the linear
layers in these models can be reduced by upto 39.58% with less than 1% increase
in perplexity. We then use LoRD to compress StarCoder 16B to 13.2B parameter
with no drop and to 12.3B with minimal drop in HumanEval Pass@1 score, in less
than 10 minutes on a single A100. The compressed models speeds up inference
by up to 22.35% with just a single line of change in code over huggingface’s implementation with pytorch backend. LoRD models remain compatible with state
of the art near-lossless quantization method such as SpQR, which allows leveraging further compression gains of quantization. Lastly, QLoRA over LoRD model
further reduces memory requirements by as much as 21.2% over vanilla QLoRA
while offering similar gains from parameter efficient fine tuning. Our work shows
LOw Rank Decomposition (LoRD) as a promising new paradigm for LLM compression. 1
1 INTRODUCTION
Code LLMs have become an integral component of Copilots that boost developer productivity (Peng
et al., 2023) and in LLM based agents (Wang et al., 2023a). These Code LLMs are as large as 34
Billion parameters for the publicly available models Roziere et al. (2023) and more than 175 Billion <br/>parameter for closed source ones Chen et al. (2021a). There is not only a pressing need for reducing<br/>model size and running models at a lower cost, but also for increasing the inference speed. The latter<br/>is especially significant for Copilot based applications.<br/>Recently, several methods have been proposed to compress and speed up inference of LLMs. Quantization (Frantar et al., 2023; Dettmers et al., 2023b) reduces the number of bits required per weight<br/>parameter of LLM by lowering the precision, and has shown significant model compression as well<br/>as speedups in low-batch decoding phases of LLMs Kim et al. (2023a). Quantization has also been<br/>shown to generalize well to quantized models Shen et al. (2023). Pruning (Sun et al., 2023a; Frantar<br/>& Alistarh, 2023) has offered another means of compression by removing connections from the neural network and hence sparsifying the weight matrices of the neural networks. Distillation Gu et al.<br/>1We will release LoRDCoder at https://huggingface.co/nolanoAI<br/>1<br/>arXiv:2309.14021v1 [cs.CL] 25 Sep 2023<br/>Preprint: Work in Progress<br/>(2023); Agarwal et al. (2023); Jung et al. (2023) method enables one to train a smaller model using<br/>a larger teacher model for supervision. While quantization and pruning methods that do not require<br/>re-training are viable means of compressing the model, distillation involves a significant amount of<br/>compute for retraining a smaller LLM, often from scratch. Here, we consider another compression<br/>paradigm of LOw Rank Decomposition (LoRD) , that does not require expensive retraining as in the<br/>case of distillation and covers up several deficiencies of the quantization and pruning compression<br/>method.<br/>Low Rank Decomposition factorizes a dense matrix of a neural network as a product of two smaller<br/>dense matrices. The LoRD model can leverage the highly optimized floating-point dense matrix<br/>multiplication kernels (NVIDIA, 2007; Blackford et al., 2002) that have been written over modern<br/>hardware. In contrast, quantized models require specialized kernels to be written, often different<br/>for each hardware backend in order to enable fast inference. Moreover, the neural network remaining fully-differentiable and all the parameters remaining trainable even after compression, unlike<br/>quantization. The LoRA Hu et al. (2022) layers of tuned models are also easier to merge back into<br/>floating point matrices compared to the quantized ones.<br/>Pruned models produce sparse matrix weights in the neural network. Matrix multiplication over<br/>sparse matrices is much slower than the resulting dense matrices in LoRD on most GPUs. Dense<br/>matrices, in addition avoid representation format overhead that sparse matrices incur from parameter<br/>reduction 2<br/>and often requires specialized kernels for reducing this overhead Dettmers et al. (2023b).<br/>Dense matrix multiplication is also easier to implement than sparse matrix multiplication, especially<br/>over quantized models.<br/>Several previous works have attempted to apply matrix decomposition methods like SVD, Tucker<br/>or Kronecker decomposition for compression (Ben Noach & Goldberg, 2020; Tahaei et al., 2022;<br/>Edalati et al., 2022). However, these have been limited to small language models like Bert (Devlin<br/>et al., 2019) and GPT2 (Radford et al., 2019), and have shown success only on narrow task-specific<br/>use cases or after retraining, often only with teacher-guided distillation supervision. These works<br/>have observed that weight matrices are not low rank and adapt methods like Singular Value Decomposition for data-aware decomposition of weights (Chen et al., 2021b; Hsu et al., 2022; Yu & Wu,<br/>2023).<br/>We, adapt these approaches for Large Language Models (Billion+ Parameters) over python code, and<br/>show that these models can be low-rank decomposed to compress and speed up inference without<br/>the need for retraining with little to no performance degradation. We study low-rank decomposition<br/>across two families of code LLMs - StarCoder and CodeGen (§2) for varying parameter sizes and<br/>establish the potential for reducing rank of models through decomposition. We then study these<br/>trends across different kinds of linear layers in a transformer block and observe the potential for<br/>upto 39.58% rank reduction with less than 1% change in perplexity.<br/>We propose various considerations for compressing the models and to achieve inference speedup<br/>on GPUs (§3.1). Using these, we achieve compression of the StarCoder 16B model offering 31.67<br/>HumanEval Chen et al. (2021a) Pass@1 score down to 13.2B parameter with similar performance of<br/>31.57 HumanEval and down to 12.3B parameter with 29.22 HumanEval score (§3.2). LoRD models,<br/>offer an inference speedup of as high as 22.35% with just one line of change in huggingface’s (§3.3).<br/>These LoRD models can be further compressed via the near-lossless quantization method of SpQR<br/>Dettmers et al. (2023b) to reduce it’s precision to 8 and 4 bits without any further reduction in<br/>HumanEval performance (§4.1). Finally, these decomposed models also reduce the memory requirements of adapter finetuning by 21.2% over QLoRA (§4.2).<br/>2This overhead in sparse matrix occurs from having to store indices/bitmasks to indicate which values are<br/>present and not. This can be very significant at low levels of sparsity. PyTorch’s sparse formats (CSR, CSC,<br/>COO) all store indices at int64 format, and for moderate levels of sparsity (&lt;50%), the sparse matrix takes up<br/>more space than a dense matrix with zero-ed out values.<br/>2<br/>Preprint: Work in Progress<br/>2 CODE LLMS ARE LOW RANK DECOMPOSABLE<br/>2.1 BACKGROUND<br/>Let an linear layer L of an LLM M with weight W ∈ R<br/>d1×d2 and bias b ∈ R<br/>d1×1<br/>. Let dmin =<br/>minimum(d1, d2) and dmax = maximum(d1, d2)<br/>A Low Rank Decomposition or Low Rank Factorization of a layer L would give us a new layer L˜<br/>with two weight matrices A ∈ R<br/>r×d2 and B ∈ R<br/>d1×r<br/>, and a bias ˜b ∈ R<br/>d1×1<br/>, where r &lt;&lt; dmin<br/>such that for a n batch of input vectors X ∈ R<br/>d2×n the batch of output vectors Y ∈ R<br/>d1×n is,<br/>Y = L˜(X) = BAX + ˜b ≈ L(X) = W X + b (1)<br/>Singular Value Decomposition (SVD) offers the best r-rank approximation of matrix W ∈ R<br/>d1×d2<br/>.<br/>First W can be decomposed as W = USV T<br/>, where U ∈ R<br/>d1×d2 and V ∈ R<br/>d2×d2 are orthogonal<br/>matrix and S ∈ R<br/>d1×d2<br/>is a diagonal matrix with entries in decreasing order. Then, by taking top-k<br/>rank, we can decompose W as a product of two low ranked matrices W ≈ BA as follows<br/>W = (U:,:rS:r,:r)<br/>| {z }<br/>B∈Rd1×r<br/>(V:r,:)<br/>| {z }<br/>A∈Rr×d2<br/>(2)<br/>where :a,:b denotes a slice operation over a matrix that gives its first a rows and b columns.<br/>Eigendecomposition is another decomposition method applicable to symmetric matrices. We can<br/>represent the eigendecomposition of a symmetric matrix W ∈ R<br/>d1×d1 as W = QΛQT<br/>. Here<br/>Q ∈ R<br/>d1×d1<br/>is an orthogonal matrix whose columns are the eigenvectors of W, and Λ ∈ R<br/>d1×d1<br/>is a diagonal matrix whose entries are the eigenvalues of W sorted in decreasing order. Similar to<br/>SVD, we can decompose W as a product of two low ranked matrices W ≈ BA by retaining only<br/>the largest r eigenvalues (and corresponding eigenvectors) as follows:<br/>W = (Q:,:rΛ:r,:r)<br/>| {z }<br/>B∈Rd1×r<br/>(Q<br/>T<br/>:r,:<br/>)<br/>| {z }<br/>A∈Rr×d1<br/>(3)<br/>Since Q is orthonormal and the eigenvalues Λ is sorted in descending order, Q:,:rQT<br/>:,:r ≈ I where I<br/>is identity matrix of dimension d1.<br/>While SVD gives the optimal low-rank decomposition of matrix, in terms of Frobenius norm, but<br/>does not take input and output data distribution into account. Approaches like weighted SVD (Hsu<br/>et al., 2022) and SVD over both weight and data (Chen et al., 2021b) have been proposed but<br/>are prohibitively expensive to scale to larger models for their requirement of backpropagation over<br/>calibration dataset. SVD over very large weight matrices is also very computationally expensive.<br/>So, we instead leverage the observation that activations in transformers are low-ranked (Feng et al.,<br/>2022) and adapt the more heuristically driven approach of Atomic Feature Mimicking (AFM) (Yu<br/>& Wu, 2023) that creates low rank matrices conditioned on a small amount of calibration data.<br/>Specifically, consider the eigen-decomposition of Covariance over Y as<br/>E[yyT<br/>] − E[y]E[y]<br/>T = QˆΛˆQˆT<br/>(4)<br/>Here Qˆ is a matrix of its eigenvectors, hence Qˆ<br/>:,:rQˆT<br/>:,:r ≈ I. Using this, we can write the output<br/>vector Y as Y ≈ Qˆ<br/>:,:rQˆT<br/>:,:rY . By writing Y in terms of W, X and b from Equation 1, we have:<br/>Y ≈ Qˆ<br/>:,:rQˆT<br/>:,:rW X + Qˆ<br/>:,:rQˆT<br/>:,:r<br/>b (5)<br/>Comparing to Equation 1, this gives us B = Qˆ<br/>:,:r ∈ R<br/>d1×r<br/>, A = QˆT<br/>:,:rW ∈ R<br/>r×d2 and ˜b =<br/>Qˆ<br/>:,:rQˆT<br/>:,:r<br/>b ≈ b. This approach is also straightforward to adapt for LLMs like LLaMa (Touvron<br/>et al., 2023), Falcon (Penedo et al., 2023), CodeLLaMa (Roziere et al., 2023) which do not have a
bias term in the linear layer by setting ˜b to zero vector.
3
Preprint: Work in Progress
2.2 EXPERIMENTAL SETTINGS
We take our python calibration dataset from the stack (Kocetkov et al., 2022) and consider the
corresponding subset of the stack smol (Bigcode, 2022) as validation data. We filter out those
sequences which are less than 1024 tokens or 10240 characters in length. We consider CodeGen and
StarCoder model family of models. CodeGen mono models are present across 350M, 2B, 6B and
16B parameters and are CodeGen models that were further trained on only python code. StarCoder
16B is the StarCoderBase 16B model further trained on only python code from the stack dataset’s
train split. We also consider StarCoderBase at 3B and 7B parameter sizes in StarCoder family due
to the lack of their monolingual counterparts. All our experiments were performed on a single A100
GPU in under an hour for each run.
For studying the trends of increase in perplexity for a reduction in rank across difference model
sizes, we set a fixed low-rank r for all the layers. Later we discuss how to achieve compression and
inference speedup via low-rank decomposition in §3
2.3 CHANGE IN PERPLEXITY ACROSS REDUCTION IN RANK
Figure 1a and 1b show the trends of increase in perplexity across reduction in rank of the weight
matrix of CodeGen and StarCoder models. For the largest models in both families, we observe only
about a 1% increase in perplexity for 10% reduction in rank, and upto 35% reduction in rank for less
than 10% increase in perplexity. The smallest model, CodeGen Mono 350M, however, can only be
decomposed to 35% rank reduction for a similar drop in perplexity. We observe that the perplexity
changes much slower for larger models as the % rank reduces, and hence can be compressed mode,
similar to observations in quantization and pruning (Li et al., 2020). It should be noted that for most
models, more than 50% leads to significant output quality degradation.
(a) Perplexity vs % Rank Reduction for CodeGen
Models.
(b) Perplexity vs % Rank Reduction for StarCoder
Models.
Figure 1: Perplexity vs %Reduction in Rank for Different Models.
3 COMPRESSION AND SPEEDUP THROUGH DECOMPOSITION
In this section, we discuss how we adapt the LOw Rank Decomposition (LoRD) for reducing the
size of model and achieving inference speedup without a significant reduction in the output quality
of the model. Following (Kim et al., 2023a), we assume memory bandwidth is the bottleneck for
inference, and thus speedups for decoding are directly proportional to the size of the transformer
model.
3.1 ACHIEVING COMPRESSION AND INFERENCE SPEEDUP
Threshold for size reduction across rank reduction: Consider a weight matrix W ∈ R
d1×d2
of a transformer layer with low rank decomposed A ∈ R
r×d2 and B ∈ R
d1×r
. The number of
parameters before and after decomposition respectively are d1d2 and r(d1 + d2). Therefore, if
4
Preprint: Work in Progress
r > d1d2
(d1+d2)
, (i.e a decomposition with small rank reduction), then the size of the model after
decomposition can even be higher than the original models. Ideally, we would want the rank r <<
d1d2
(d1+d2)
or r << dmin.
Matrix Aspect Ratio and Compression: Let the ratio of the smaller dimension to the larger dimension of the matrix (i.e. the aspect ratio) be α =
dmin
dmax
. For square matrix, α = 1 and for tall or
fat matrices α << 1. We can rewrite, the percentage change in parameters from decomposition, in
terms of percent change in rank %∆r = 100 ∗
dmin−r
dmin
% and aspect ratio as:
100 ∗
r(dmax + dmin) − dmaxdmin
dmaxdmin
= 100α − (1 + α)%∆r (6)
It should be noted that change in parameters from decomposition can either be positive (the number
of parameters increased after decomposition), or negative (the number of parameters decreased after
decomposition). In order to achieve model compression and consequently inference speedups, one
would want a very high negative percentage change in parameters.
Figure 2: Parity Point across various aspect
ratios (α) of the different linear layers in
transformers.
Parity Point for Compression across Rank Reduction: Using Eq. 6, one can observe that little reduction in rank, may lead to increase in model parameters instead of decreasing. For instance, square matrices (α = 1) will have 100% increase (i.e doubling
in size), then %∆r → 0+ and only after the rank
is reduced by more than 50%, will the Parity Point
of the rank reduction be reached, that offers same or
lesser number of a parameter in the decomposed layer
as the original matrix. This parity point for tall or fat
matrices (α → 0+), can be achieved with a very small
percent reduction in rank and can start giving a reduction in model size. For compression to be achieved,
we would want to reduce the rank by an amount to
cross this parity point threshold. However, reducing
the rank by a lot can degrade performance significantly. So we must take the aspect ratio into account,
in order to achieve compression without much reduction in rank (and hence no significant degradation in
output quality)
A transformer model had different aspect ratios across its various linear layers, α = 1.00 for output
projection after attention, α = 0.96 for Multi-query attention (Shazeer, 2019) projections, α = 0.25
for typical MLP projections with intermediate expansion factor of 4 as in the original transformer
and as low as α = 0.12 for the embedding and language model head projection of CodeGen 16B
with 51200 vocab size. Figure 2 plots the % change in the size of the model across % reduction
in rank for matrices with different aspect ratios. For square matrices and near square matrices,
a small rank reduction doubles the size of the linear layer after decomposition, and only after its
parity point of 50% reduction is the size after decomposition, the same as original matrix. By this
extent of rank decomposition, the performance starts to significantly degrade, as seen in §2.3. All the
previous works on smaller models, address this by retraining the model (Yu & Wu, 2023; Chen et al.,
2021b; Hsu et al., 2022; Ben Noach & Goldberg, 2020), often via knowledge distillation supervision
(Hinton et al., 2015; Sanh et al., 2019) on specific narrow tasks. However, retraining is infeasible
for larger models. Thus, we skip matrices with very high aspect ratios such as output projection or
multi-query attention for decomposition. In contrast, the weights in MLP achieve parity at only 20%
rank reduction. While embeddings and LM Head can be compressed through decomposition, as they
have been for smaller transformer models (Baevski & Auli, 2019; Lan et al., 2020), they contribute
only a very small portion of the weight of the model. So, we do not consider decomposing these
matrices. In order to reduce the aspect ratio of matrices, we group layers with the same input vector
to have the same bottleneck matrix after decomposition. Doing so, enables re-use of computation,
and sharing of weights, as well as bringing the aspect ratio down to achieve compression as lower
rank reduction. Candidate layers for grouping include the query, key and value projection matrices
5
Preprint: Work in Progress
in multi-headed attention with aspect ratio reduced to α = 0.33 and the gating layer in SwiGLU
(Shazeer, 2020) with first linear linear of MLP in models like LLaMa (Touvron et al., 2023) with
α = 0.1875.
(a) CodeGen 16B. (b) StarCoder 16B.
Figure 3: Parameter Reduction vs perplexity for decomposition across various layers.
Trends across different layers in a transformer block: In addition to considering the parity point
into account for deciding which layers to decompose, we also additionally study the sensitivity of
each of these layers to low rank decomposition across the large model in the two model families.
Figure 3 shows the increase in perplexity vs reduction in model parameters for the two models. For
both models, decomposing all the linear layers achieves the parity point much later than any one of
these linear layers with low aspect ratio. For CodeGen, the attention weight matrix (query, key and
values projection) offers least increase in perplexity for the biggest drop in parameter count, make
this layer the most suitable candidate to be decomposed. It shows less than 1% increase in perplexity
even after 39.58% rank reduction. We observe the mlp 2 (downscaling mlp) to be a better candidate
for decomposition than mlp 1 (upscaling mlp) across both models. This makes mlp 2 to be a good
candidate for low-rank decomposition over the StarCoder model.
Hardware Considerations: On modern hardware accelerators like GPU and their corresponding
software stack, matrix multiplication kernels are faster if their dimensions are divisible by a high
factor of 2. So, we consider ranks at a reduction of approximately every 10%, rounded off to the
nearest multiple of 128 in our experiments.
3.2 PERFORMANCE OF COMPRESSED MODELS
We consider the largest models of StarCoder and CodeGen family (16B) and perform low-rank decomposition on both with varying ranks. We consider decomposing layers that offers most parameter
reduction (§3.1) with least increase in perplexity - mlp 2 for StarCoder and attention for CodeGen.
We report the Pass@1 and Pass@10 scores over the Human Eval dataset (Chen et al., 2021a) using the code-eval GitHub repo (Bacaj, 2023) in Table 1. We observe that StarCoder models can be
low rank decomposed to 13.2B parameters (50% rank reduction) with no drop in Pass@1 performance and upto 12.3B parameters (62.5% rank reduction) with very little drop. CodeGen models
shows similar trend in drop in Human Eval performance when measured in terms of rank reduction.
However, in terms of parameter reduction count, while showing very little perplexity change with
large reduction in rank (Fig. 3a), shows much more drop in its HumanEval score when measured
in terms of parameter count reduction due to a higher aspect ratio of the matrix being decomposed.
It should be noted that for certain compressed models, the Pass@1 even slightly improves over the
base model. Similar trend of slight improvements from compression across various metrics and
benchmarks has been observed in the case of other compression attempts (Frantar & Alistarh, 2023;
Cerebras, 2022).
6
Preprint: Work in Progress
Starcoder 16B CodeGen 16B Mono
Model Type Rank HumanEval Score Model Type Rank HumanEval Score
Pass @ 1 Pass @ 10 Pass @ 1 Pass @ 10
Base Model 6144 31.67 48.28 Base Model 6144 29.02 46.34
LoRDCoder 14.9B 4480 33.18 48.41 LoRDCoder 15.9B 4480 29.08 46.95
LoRDCoder 14.5B 4096 31.69 45.12 LoRDCoder 15.6B 4096 28.90 46.24
LoRDCoder 13.8B 3584 30.90 47.56 LoRDCoder 15.1B 3584 28.54 45.73
LoRDCoder 13.2B 3072 31.57 45.36 LoRDCoder 14.7B 3072 27.99 43.29
LoRDCoder 12.6B 2560 29.84 42.31 LoRDCoder 14.3B 2560 27.32 45.12
LoRDCoder 12.3B 2304 29.22 40.12 LoRDCoder 14.1B 2304 27.07 41.46
Table 1: Human Eval Score of LoRD across StarCoder and CodeGen.
3.3 SPEEDUP FROM LORD
We next consider accessing the inference speedup (forward pass) of the models over the standard
cuBLAS floating point kernels. We consider the standard Huggingface implementation (Wolf et al.,
2020) of Starcoder with pytorch backend (Paszke et al., 2019) utilizing standard cuBLAS kernels
on A100 GPUs. LoRD decomposed models were implemented by modifying just one line of code
to replace an MLP with an extra linear layer 3
. We benchmark over 1024 tokens and 512 tokens
sequence, averaged across 10 runs with warm up of 3 runs. We plot relative time taken and model
size across reduction in rank in Figure
Figure 4: Time and Model size of StarCoder 16B across ranks.
4.
Inference speedups as high as 22.35% are observed
for decomposed models. The lines in the graph are
generally downward sloping, Therefore reduction in
rank beyond 25% generally implies less inference time
and reduction in model size. However, the underlying
hardware (and pertaining software kernels) also significantly affect the speedup gains. We notice huge gains,
whenever the rank is rounded off to a multiple of a very
high power of 2 (like 4096 and 2560 at 33% and 58%
rank reduction), despite very little reduction in model
size. In contrast, for certain ranks which are multiples
of a lesser power of 2 (like 3584 and 2304 at 41% and
62% rank reduction) are slower than those at slightly
higher ranks. It is worth noting that affect of hardware
inefficient matrix shape is less significant for longer tokens sequence of 1024 because the O(n
2
) attention overhead starts becoming more significant,
especially in the absence of SoTA attention implementation techniques (Rabe & Staats, 2021; Dao
et al., 2022; Dao, 2023) as in the case of Huggingface’s implementations.
4 COMBINING LORD WITH QUANTIZATION AND LORA
4.1 QUANTIZATION
While LoRD enables compression at same precision level, we study whether the decomposed models can be further compressing through quantization. Table 2 shows the HumanEval pass@1 results
for the different LoRDCoder across 8 and 4 bit quantization levels, using the near lossless quantization technique of SpQR (Dettmers et al., 2023b). We observe that the LoRD models can be
combined with quantization for further compression, showing no performance drop for 8-bit and
very little performance drop on 4-bit quantization for most models. Slight increase in HumanEval
after quantization is also observed, similar to Pangu-Coder2 (Shen et al., 2023).
4.2 PARAMETER EFFICIENT TUNING OF LORD MODELS
3
nn.Linear(in, out) -> nn.Sequential(nn.Linear(in, rank), nn.Linear(rank, out))
7
Preprint: Work in Progress
Model Pass@1@FP16 Pass@1@8-bit Pass@1@4-bit
LoRDCoder 14.9B 33.18 33.17 32.01
LoRDCoder 14.5B 31.69 31.58 32.74
LoRDCoder 13.8B 30.90 31.10 30.73
LoRDCoder 13.2B 31.57 31.52 32.01
LoRDCoder 12.6B 29.84 29.87 30.22
LoRDCoder 12.3B 29.22 29.14 29.45
Table 2: Human Eval score of quantized LoRDCoder models.
d x d
d x r’
r’ x d
d x r’
r’ x d
d x r
r x d
d
r’
h
x
d
r’
h
x
r
LoRD
Weights
a) LoRA b) LoRD + LoRA
Pr e t r a i n ed
W e i g h t s
A R
X R
Y R
W R
X R
Y R B R
Figure 5: LoRA vs LoRD + LoRA.
We next test the potential for using
LoRD to further reduce the memory usage over existing parameterefficient techniques. We consider the
code instruction dataset (Chaudhary,
2023) and filter those examples that
pertains to python programming language. We use QLoRA (Dettmers
et al., 2023a), which is an even more
memory efficient version of LoRA
(Hu et al., 2022) storing the weights
in quantized format, for fine-tuning
for 1 epoch. We compare results
from fine-tuning two of the decomposed models LoRDCoder 13.2B and
LoRDCoder 12.3B model to the StarCoder model. We observe a HumanEval pass@1 of 37.80 and 37.62 across LoRDCoder 13.2B
and LoRDCoder 12.3B fine-tuning, competitive to the performance of 37.74 offered by StarCoder
model.
5 RELATED WORK
There is a growing interest in compressing pretrained Large Language Models. Several recent attempts have been dedicated to the quantization of weights of LLMs (Frantar et al., 2023; Lin et al.,
2023; Yuan et al., 2023; Park et al., 2022; Kim et al., 2023b; Chee et al., 2023; Li et al., 2023a) with
tricks such as outlier separation (Dettmers et al., 2022; Dettmers & Zettlemoyer, 2022; Dettmers
et al., 2023c; Wei et al., 2022; Kim et al., 2023a; Lee et al., 2023). Some attempts also quantize
the activations (intermediate representations) in addition to weights to speed up computation time
(Shao et al., 2023; Xiao et al., 2023). The works in quantization that are closest to us is the LowRank Compensation (LoRC) Strategy (Yao et al., 2023; Wu et al., 2023), where the difference of
the quantized matrix to the original matrix is approximated by a product of low-rank matrices. Our
work decomposes the entire matrix for compression.
Pruning neural networks Liang et al. (2021), unlike quantization, reduces the number of parameters
in a model by removing unimportant weights or connections. Several techniques have been proposed
to scale pruning methods for LLMs (Sun et al., 2023a; Frantar & Alistarh, 2023; Ma et al., 2023).
However, pruning as a means of compression is yet to become viable due to no speedups over sparse
matrices without significant performance drop at extreme levels of sparsity or structured sparsity
(Zhu et al., 2023). With low-rank decomposition, we propose an alternate method for reducing
model parameters that offer speedup even at a little reduction in parameter count. Certain works
have also attempted to (Ren & Zhu, 2023; Li et al., 2023b) to split a dense matrix as a sum of lowrank matrices and a sparse matrix. However, these methods require retraining and have been shown
to work only for Language Models of less than a billion parameters.
Low rank decomposition has been proposed for smaller language models like Bert or GPT2 before
using SVD decomposition (Ben Noach & Goldberg, 2020) and Kronecker decompositions (Tahaei
et al., 2022; Edalati et al., 2022). Hsu et al. (2022) modified SVD to be data aware based on approximate second-order gradient information. A better weighted SVD was proposed by (Hua et al.,
2022). Chen et al. (2021b) proposed a data aware decomposition method with a provably optimal
closed-form solution, utilizing a large amount of data points over specific tasks to decompose. Sev8
Preprint: Work in Progress
eral recent works (Yu & Wu, 2023; Feng et al., 2022) have shown that while the weight matrix of
neural networks is not inherently low-rank, the intermediate representations are, thus propose to
decompose based on representations. All these works have focused on small language models and
require re-training. We proposed low-rank decomposition for compressing neural networks without
the need for retraining. The factorization has also been used just for the embedding layers (Baevski
& Auli, 2019; Lan et al., 2020), as they are good candidates due to their very low aspect ratio of
0.015, where a reduction of rank by even 5% would lead to reduction in number of parameters after
decomposition.
There is also a growing interest in fine-tuning large language models Taori et al. (2023); Chiang
et al. (2023); Wang et al. (2023b); Sun et al. (2023b). With the large memory requirements for finetuning full parameters of the LLM, the more parameter-efficient fine-tuning methods like LoRA
(Hu et al., 2022) are getting widely adopted. These methods freeze the original LLM weights, and
attach two low-rank matrices or adapters, in a skip-connection (He et al., 2016) to the linear layers
of the model. These parameter-efficient fine-tuning approaches have seen improvements in lower
activation memory (Zhang et al., 2023) or by keeping non-trainable model weights at 4-bit precision
(Dettmers et al., 2023a). Our work, while focused on compression through low-rank decomposition,
can also enable more efficient fine-tuning, especially in conjunction with existing methods.
6 CONCLUSION
We studied the compression of monolingual code generation models through a novel one-shot compression paradigm of low-rank decomposition. We analyse the change in perplexity with change
in rank across the model families of StarCoder and CodeGen as well as their individual layers and
observe that the rank of these models can be reduced by upto 39.58% with less than 1% change in
perplexity. We then proposed considerations for one-shot compressing these models through LOw
Rank Decomposition (LoRD) in under 10 minutes. Consequently, we compress StarCoder 16B
to 13.2B with no drop in HumanEval pass@1 and very little drop in HumanEval pass@1 to 12.3B
parameters. With a minimal change in code over huggingface’s default inference code of just one
line, we gain speedups of up to 22.35%. The LoRD models are also compatible with near lossless
quantization techniques of SpQR, which offers gains of quantization based compression in addition
to ones from decomposition. The LoRD models also reduce memory requirements by as much as
21.2% over vanilla QLoRA fine-tuning.
7 BROADER IMPACT AND FUTURE WORK
Our work on LoRD, compresses code LLMs which enables them to run on smaller GPUs including
as consumer grade GPUs. This is especially of pressing importance for the next few years when
the shortage of GPU supply is relative to the increasing demand in today’s market. Moreover, faster
inference helps reduce the GPU cycles, enabling lower running costs and lower power consumption
for LLM inference. Our work helps reduce the carbon emissions incurred and moves towards a
greener NLP. Through compression, our work also promotes inference at the edge, and therefore
opening room for applications involving strict privacy requirements. Lower latency will also help
improve the User Experience in applications like CoPilots where lag between suggestions can impact developer’s productivity. Several of these benefits of LoRD such as lower cost and energy
consumption are also applicable for fine-tuning use cases of LLMs.
Our work opens up a new paradigm for compression via Low Rank Decomposition over Large
Language Models in a single shot without the need for retraining. Since, LoRD models can leverage
existing floating point kernels across BLAS and cuBLAS, in contrast to quantization, these are
much easier to implement and reap inference benefits. Our study on hardware considerations for
speedup also opens up the potential for tuning the rank of decomposed models to fit best on the
target hardware and the accompanying GEMM kernels. While our study is limited to monolingual
code LLMs, the low rank decomposition technique is general and not specific to code domain. Thus
exploring its applicability to more general purpose models like LLaMa is a promising direction
for the compression of transformer LLMs beyond quantization. Another interesting unexplored
question is whether the LoRA or QLoRA modules fine-tuned on original models, can be plugged in
as-is for the LoRD models without any performance drop””
Alright, let’s simplify this down to the basics. Imagine you have a giant room full of filing cabinets, and each cabinet represents a piece of information needed for a computer program to understand and produce code, like creating software. This room is like a model that helps developers by suggesting or writing code. As you can imagine, a room that big takes a lot of time and energy to search through, not to mention the space it occupies!

Now, suppose we find a clever way to organize this room. Instead of having all these separate, bulky filing cabinets, we rearrange them into fewer, smaller cabinets but in such a way that we can still find all the information we need, just as quickly – if not faster. This rearrangement process is kind of what Ayush Kaushal and their team have done with computer programs used in coding, specifically, a kind called Large Language Models (LLMs) that are trained to understand and generate code.

Their method is called “LOw Rank Decomposition” (LoRD), which is a fancy way of saying they found a smart way to compress or shrink down the model without losing its effectiveness. Just like how rearranging the filing room makes it easier, quicker, and cheaper to find the files you need, LoRD makes these LLMs faster, smaller, and less expensive to run.

Why does this matter? Well, smaller and faster models mean they can work on less powerful computers, making this technology accessible to more people. It also means they can operate more efficiently, saving on electricity and contributing less to carbon emissions, which is better for our planet.

Moreover, because these compressed models still understand and produce code just as well, developers can use them without having to go through a lengthy process of teaching the model all over again, which saves a tremendous amount of time and resources. They even showed that we could squeeze these models down pretty significantly – making them up to 39.58% smaller! – with little to no loss in their ability to generate code effectively.

Finally, the team found these slimmed-down models could work even better when combined with other techniques (quantization and fine-tuning methods), making them even more efficient and adaptable.

In essence, this research opens the door to making powerful coding-assist tools more accessible and environmentally friendly, bringing the benefits of AI to more people while conserving valuable resources.
I stumbled across a Github repo for “Low-Rank adapter extraction for fine-tuned transformers model. This repository contains code for extracting LoRA adapters from finetuned transformers models, using Singular Value Decomposition (SVD).

LoRA (Low-Rank Adaptation) is a technique for parameter-efficient fine-tuning of large language models. The technique presented here allows extracting PEFT compatible Low-Rank adapters from full fine-tunes or merged model.

A bit more about PEFT. PEFT (Parameter-Efficient Fine-Tuning) is a library for efficiently adapting large pretrained models to various downstream applications without fine-tuning all of a model’s parameters because it is prohibitively costly. PEFT methods only fine-tune a small number of (extra) model parameters - significantly decreasing computational and storage costs - while yielding performance comparable to a fully fine-tuned model. This makes it more accessible to train and store large language models (LLMs) on consumer hardware.

PEFT is integrated with the Transformers, Diffusers, and Accelerate libraries to provide a faster and easier way to load, train, and use large models for inference.

Now, I found a huggingface profile that has a Low-rank decomposition of ezelikman’s (original paper author) quietstar-8-ahead model using mistralai/Mistral-7B-v0.1 as a base. This was created using LoRD. This is the folder contents: “”
Upload folder using huggingface_hub
b17482f
VERIFIED
2 days ago
.gitattributes
1.52 kB
initial commit
2 days ago
README.md
391 Bytes
Upload folder using huggingface_hub
2 days ago
adapter_config.json
711 Bytes
Upload folder using huggingface_hub
2 days ago
adapter_model.safetensors
336 MB
LFS
Upload folder using huggingface_hub
2 days ago”“

How can i utilize the contents in this folder for a model? I have a potential guide located here: ““LoRA methods
A popular way to efficiently train large models is to insert (typically in the attention blocks) smaller trainable matrices that are a low-rank decomposition of the delta weight matrix to be learnt during finetuning. The pretrained model’s original weight matrix is frozen and only the smaller matrices are updated during training. This reduces the number of trainable parameters, reducing memory usage and training time which can be very expensive for large models.

There are several different ways to express the weight matrix as a low-rank decomposition, but Low-Rank Adaptation (LoRA) is the most common method. The PEFT library supports several other LoRA variants, such as Low-Rank Hadamard Product (LoHa), Low-Rank Kronecker Product (LoKr), and Adaptive Low-Rank Adaptation (AdaLoRA). You can learn more about how these methods work conceptually in the Adapters guide. If you’re interested in applying these methods to other tasks and use cases like semantic segmentation, token classification, take a look at our notebook collection!

This guide will show you how to quickly train an image classification model - with a low-rank decomposition method - to identify the class of food shown in an image.

Some familiarity with the general process of training an image classification model would be really helpful and allow you to focus on the low-rank decomposition methods. If you’re new, we recommend taking a look at the Image classification guide first from the Transformers documentation. When you’re ready, come back and see how easy it is to drop PEFT in to your training!

Before you begin, make sure you have all the necessary libraries installed.

Copied
pip install -q peft transformers datasets
Dataset
In this guide, you’ll use the Food-101 dataset which contains images of 101 food classes (take a look at the dataset viewer to get a better idea of what the dataset looks like).

Load the dataset with the load_dataset function.

Copied
from datasets import load_dataset

ds = load_dataset(“food101”)
Each food class is labeled with an integer, so to make it easier to understand what these integers represent, you’ll create a label2id and id2label dictionary to map the integer to its class label.

Copied
labels = ds[“train”].features[“label”].names
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
label2id[label] = i
id2label[i] = label

id2label[2]
“baklava”
Load an image processor to properly resize and normalize the pixel values of the training and evaluation images.

Copied
from transformers import AutoImageProcessor

image_processor = AutoImageProcessor.from_pretrained(“google/vit-base-patch16-224-in21k”)
You can also use the image processor to prepare some transformation functions for data augmentation and pixel scaling.

Copied
from torchvision.transforms import (
CenterCrop,
Compose,
Normalize,
RandomHorizontalFlip,
RandomResizedCrop,
Resize,
ToTensor,
)

normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
train_transforms = Compose(
[
RandomResizedCrop(image_processor.size[“height”]),
RandomHorizontalFlip(),
ToTensor(),
normalize,
]
)

val_transforms = Compose(
[
Resize(image_processor.size[“height”]),
CenterCrop(image_processor.size[“height”]),
ToTensor(),
normalize,
]
)

def preprocess_train(example_batch):
example_batch[“pixel_values”] = [train_transforms(image.convert(“RGB”)) for image in example_batch[“image”]]
return example_batch

def preprocess_val(example_batch):
example_batch[“pixel_values”] = [val_transforms(image.convert(“RGB”)) for image in example_batch[“image”]]
return example_batch
Define the training and validation datasets, and use the set_transform function to apply the transformations on-the-fly.

Copied
train_ds = ds[“train”]
val_ds = ds[“validation”]

train_ds.set_transform(preprocess_train)
val_ds.set_transform(preprocess_val)
Finally, you’ll need a data collator to create a batch of training and evaluation data and convert the labels to torch.tensor objects.

Copied
import torch

def collate_fn(examples):
pixel_values = torch.stack([example[“pixel_values”] for example in examples])
labels = torch.tensor([example[“label”] for example in examples])
return {“pixel_values”: pixel_values, “labels”: labels}
Model
Now let’s load a pretrained model to use as the base model. This guide uses the google/vit-base-patch16-224-in21k model, but you can use any image classification model you want. Pass the label2id and id2label dictionaries to the model so it knows how to map the integer labels to their class labels, and you can optionally pass the ignore_mismatched_sizes=True parameter if you’re finetuning a checkpoint that has already been finetuned.

Copied
from transformers import AutoModelForImageClassification, TrainingArguments, Trainer

model = AutoModelForImageClassification.from_pretrained(
“google/vit-base-patch16-224-in21k”,
label2id=label2id,
id2label=id2label,
ignore_mismatched_sizes=True,
)
PEFT configuration and model
Every PEFT method requires a configuration that holds all the parameters specifying how the PEFT method should be applied. Once the configuration is setup, pass it to the get_peft_model() function along with the base model to create a trainable PeftModel.

Call the print_trainable_parameters() method to compare the number of parameters of PeftModel versus the number of parameters in the base model!

LoRA decomposes the weight update matrix into two smaller matrices. The size of these low-rank matrices is determined by its rank or r. A higher rank means the model has more parameters to train, but it also means the model has more learning capacity. You’ll also want to specify the target_modules which determine where the smaller matrices are inserted. For this guide, you’ll target the query and value matrices of the attention blocks. Other important parameters to set are lora_alpha (scaling factor), bias (whether none, all or only the LoRA bias parameters should be trained), and modules_to_save (the modules apart from the LoRA layers to be trained and saved). All of these parameters - and more - are found in the LoraConfig.

Copied
from peft import LoraConfig, get_peft_model

config = LoraConfig(
r=16,
lora_alpha=16,
target_modules=[“query”, “value”],
lora_dropout=0.1,
bias=“none”,
modules_to_save=[“classifier”],
)
model = get_peft_model(model, config)
model.print_trainable_parameters()
“trainable params: 667,493 || all params: 86,543,818 || trainable%: 0.7712775047664294”
Training
For training, let’s use the Trainer class from Transformers. The Trainer contains a PyTorch training loop, and when you’re ready, call train to start training. To customize the training run, configure the training hyperparameters in the TrainingArguments class. With LoRA-like methods, you can afford to use a higher batch size and learning rate.

AdaLoRA has an update_and_allocate() method that should be called at each training step to update the parameter budget and mask, otherwise the adaptation step is not performed. This requires writing a custom training loop or subclassing the Trainer to incorporate this method. As an example, take a look at this custom training loop.

Copied
from transformers import TrainingArguments, Trainer

account = “stevhliu”
peft_model_id = f”{account}/google/vit-base-patch16-224-in21k-lora”
batch_size = 128

args = TrainingArguments(
peft_model_id,
remove_unused_columns=False,
evaluation_strategy=“epoch”,
save_strategy=“epoch”,
learning_rate=5e-3,
per_device_train_batch_size=batch_size,
gradient_accumulation_steps=4,
per_device_eval_batch_size=batch_size,
fp16=True,
num_train_epochs=5,
logging_steps=10,
load_best_model_at_end=True,
label_names=[“labels”],
)
Begin training with train.

Copied
trainer = Trainer(
model,
args,
train_dataset=train_ds,
eval_dataset=val_ds,
tokenizer=image_processor,
data_collator=collate_fn,
)
trainer.train()
Share your model
Once training is complete, you can upload your model to the Hub with the push_to_hub method. You’ll need to login to your Hugging Face account first and enter your token when prompted.

Copied
from huggingface_hub import notebook_login

notebook_login()
Call push_to_hub to save your model to your repositoy.

Copied
model.push_to_hub(peft_model_id)
Inference
Let’s load the model from the Hub and test it out on a food image.

Copied
from peft import PeftConfig, PeftModel
from transfomers import AutoImageProcessor
from PIL import Image
import requests

config = PeftConfig.from_pretrained(“stevhliu/vit-base-patch16-224-in21k-lora”)
model = AutoModelForImageClassification.from_pretrained(
config.base_model_name_or_path,
label2id=label2id,
id2label=id2label,
ignore_mismatched_sizes=True,
)
model = PeftModel.from_pretrained(model, “stevhliu/vit-base-patch16-224-in21k-lora”)

url = “https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/beignets.jpeg”
image = Image.open(requests.get(url, stream=True).raw)
image

Convert the image to RGB and return the underlying PyTorch tensors.

Copied
encoding = image_processor(image.convert(“RGB”), return_tensors=“pt”)
Now run the model and return the predicted class!

Copied
with torch.no_grad():
outputs = model(**encoding)
logits = outputs.logits

predicted_class_idx = logits.argmax(-1).item()
print(“Predicted class:”, model.config.id2label[predicted_class_idx])
“Predicted class: beignets”
←
Prompt-based methods
IA3
→”"
To utilize the contents of the folder you found on GitHub for a model using LoRA (Low-Rank Adaptation) and integrate it with Parameter-Efficient Fine-Tuning (PEFT), you’re actually looking to take a sophisticated approach to employing machine learning models, specifically for image classification in your stated goal. The process can be outlined in a few structured steps, adapted from your description:

### Step 1: Environment Preparation
Ensure your environment is ready by installing necessary libraries:

pip install -q peft transformers datasets


### Step 2: Dataset Preparation
For your task, you’re focusing on an image classification model. Load the specified dataset:

from datasets import load_dataset
ds = load_dataset(“food101”)


Set up label dictionaries for easy reference:

labels = ds[“train”].features[“label”].names
label2id, id2label = {label: i for i, label in enumerate(labels)}, {i: label for i, label in enumerate(labels)}


### Step 3: Image Processor Setup
Prepare your image data correctly by setting up an image processor:

from transformers import AutoImageProcessor
image_processor = AutoImageProcessor.from_pretrained(“google/vit-base-patch16-224-in21k”)


Implement the necessary transformations:

from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, RandomResizedCrop, RandomHorizontalFlip
# Define your train and validation transformations here based on the example provided in the guide.


Apply transformations to your dataset:

# Apply preprocess_train and preprocess_val functions on your train and validation datasets.


### Step 4: Model Loading and PEFT Configuration
Load a pre-trained model that will serve as your base for fine-tuning:

from transformers import AutoModelForImageClassification
model = AutoModelForImageClassification.from_pretrained(
    “google/vit-base-patch16-224-in21k”,
    label2id=label2id,
    id2label=id2label
)


Configure your LoRA settings:

from peft import LoraConfig, get_peft_model
config = LoraConfig(
    r=16,  # Rank of the decomposition
    target_modules=[“query”, “value”],  # Specifying where the LoRA adapters will be inserted.
    lora_alpha=16,
    bias=“none”,
    modules_to_save=[“classifier”]  # Specify if any other module needs to be saved along with LoRA adapters.
)
model = get_peft_model(model, config)


### Step 5: Training
Set up your training arguments:

from transformers import TrainingArguments, Trainer
args = TrainingArguments(
    # Define your training arguments here.
)


Initiate the training with the Trainer class:

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    data_collator=collate_fn,
)
trainer.train()


### Step 6: Model Sharing
To share your fine-tuned model, first log in to Hugging Face:

from huggingface_hub import notebook_login
notebook_login()


Push your model to the Hub:

model.push_to_hub(“your-model-name-here”)


### Step 7: Inference
For performing inference, ensure you load the correct model and configuration from the Hub. Then, pass an image through the model to get predictions.

Ensure you adapt each step’s example codes to your specific scenario, especially modifying the paths and model names as necessary. This comprehensive approach will allow you to fine-tune a model on the Food-101 dataset with LoRA and possibly other PEFT techniques, share it and perform inference efficiently.


------------------------------------------------------------------------------------------------------

###  New conversation :

I am actually trying to use the LORD repo to understand how i can implement the main idea from this paper: "" Quiet-STaR: Language Models Can Teach Themselves to
Think Before Speaking
Eric Zelikman
Stanford University
Georges Harik
Notbad AI Inc
Yijia Shao
Stanford University
Varuna Jayasiri
Notbad AI Inc
Nick Haber
Stanford University
Noah D. Goodman
Stanford University
Abstract
When writing and talking, people sometimes pause to think. Although
reasoning-focused works have often framed reasoning as a method of
answering questions or completing agentic tasks, reasoning is implicit
in almost all written text. For example, this applies to the steps not
stated between the lines of a proof or to the theory of mind underlying
a conversation. In the Self-Taught Reasoner (STaR, Zelikman et al. 2022),
useful thinking is learned by inferring rationales from few-shot examples
in question-answering and learning from those that lead to a correct
answer. This is a highly constrained setting – ideally, a language model
could instead learn to infer unstated rationales in arbitrary text. We
present Quiet-STaR, a generalization of STaR in which LMs learn to
generate rationales at each token to explain future text, improving their
predictions. We address key challenges, including 1) the computational cost
of generating continuations, 2) the fact that the LM does not initially know
how to generate or use internal thoughts, and 3) the need to predict beyond
individual next tokens. To resolve these, we propose a tokenwise parallel
sampling algorithm, using learnable tokens indicating a thought’s start and
end, and an extended teacher-forcing technique. Encouragingly, generated
rationales disproportionately help model difficult-to-predict tokens and
improve the LM’s ability to directly answer difficult questions. In particular,
after continued pretraining of an LM on a corpus of internet text with
Quiet-STaR, we find zero-shot improvements on GSM8K (5.9%→10.9%)
and CommonsenseQA (36.3%→47.2%) and observe a perplexity improvement of difficult tokens in natural text. Crucially, these improvements
require no fine-tuning on these tasks. Quiet-STaR marks a step towards
LMs that can learn to reason in a more general and scalable way.
“Life can only be understood backwards; but it must be lived forwards.”
— Sren Kierkegaard
1 Introduction
Much of the meaning of text is hidden between the lines: without understanding why
statements appear in a document, a reader has only a shallow understanding. Moreover,
this has been repeatedly shown to be true for LMs as well, in the contexts of tasks ranging
from commonsense reasoning to theorem proving to programming (Wei et al., 2022b; Nye
et al., 2021; Zelikman et al., 2022; 2023a; Kojima et al., 2022). Reasoning about implications
of text to predict later text has consistently been shown to improve LM performance on a
variety of tasks, but methods for allowing LMs to learn from their reasoning (e.g., Zelikman
et al. 2022) have focused on solving individual tasks or predefined sets of tasks (e.g., Wei
et al. 2021b). These works rely on carefully curated datasets to provide either specific
reasoning tasks or in some cases, the reasoning itself. We instead ask, if reasoning is implicit
in all text, why shouldn’t we leverage the task of language modeling to teach reasoning?
1
arXiv:2403.09629v2 [cs.CL] 18 Mar 2024
START +
Sampled
Thought 4 + 1 END
LM LM LM
1 Discard
2 + 1 2
START +
1) Think
= 5 4 \n
2 2 END
Thought1 Thought3
Thought4 Thought5
Thought7
+
LM
LM LM LM
LM
Reward
Harms
Prediction
Helps
Prediction
LM
Thought2
LM LM LM LM LM
4
Thought0
LM
2) Talk 3) Learn
Original Text
Sampled
Thought
Thought6
Figure 1: Quiet-STaR. We visualize the algorithm as applied during training to a single
thought. We generate thoughts, in parallel, following all tokens in the text (think). The
model produces a mixture of its next-token predictions with and without a thought (talk).
We apply REINFORCE, as in STaR, to increase the likelihood of thoughts that help the model
predict future text while discarding thoughts that make the future text less likely (learn).
In particular, the Self-Taught Reasoner (STaR, Zelikman et al. 2022) showed that LMs
can bootstrap their reasoning ability on question-answering (QA) datasets by sampling
rationales to attempt to answer questions, training on rationales if they led to a correct final
answer, and then repeating this to iteratively solve more difficult problems. Yet, training
from curated QA datasets limits the scale and generalizability of the rationales. QA datasets,
especially high-quality ones, require thoughtful curation and will inherently only ever cover
a subset of reasoning tasks. Thus, we extend STaR – instead of the LM learning to reason
on particular tasks like mathematical QA, we train an LM to generate reasoning that helps
it infer future text from a large internet text corpus. As a result, we allow the LM to learn
from the diverse tasks present in language (Weber et al., 2021). This builds on an intuition
essential to the current language modeling paradigm, namely, that ”language models are
unsupervised multitask learners” (Radford et al., 2019). Thus, as in STaR, we leverage the
LM’s pre-existing reasoning ability to generate rationales and train the LM on them with a
REINFORCE-based reward (Williams, 1992). We refer to this technique as Quiet-STaR, as it
can be understood as applying STaR “quietly”, training the model to think before it speaks.
Broadly, Quiet-STaR proceeds by generating rationales after every token to explain future
text (think), mixing the future-text predictions with and without rationales (talk), and then
learning to generate better rationales using REINFORCE (learn). We apply Quiet-STaR
to Mistral 7B (Jiang et al., 2023) using the web text datasets OpenWebMath (Paster et al.,
2023) and Colossal Clean Crawled Corpus (C4, Raffel et al. 2020). We find that, even
without dataset-specific fine-tuning, Quiet-STaR results in improvements to zero-shot directreasoning abilities on CommonsenseQA (36.3%→47.2%) and GSM8K (5.9%→10.9%), and
that these improvements consistently increase with the number of tokens used in the LM’s
internal thoughts. Lastly, we qualitatively investigate patterns in the generated rationales.
In solving this task, we make the following contributions:
1. We generalize STaR to learn reasoning from diverse unstructured text data. To our
knowledge, this is the first work explicitly training LMs to reason generally from
text, rather than on curated reasoning tasks or collections of reasoning tasks.
2. We propose and implement a parallel sampling algorithm that makes our training
procedure scalable, generating rationales from all token positions in a given string.
3. We introduce custom meta-tokens at the start and end of each thought to allow
the LM to learn that it should be generating a rationale and when it should make a
prediction based on that rationale.
4. We apply a mixing head to retrospectively determine how much to incorporate the
next-token prediction from a given thought into the current next-token prediction.
5. We show that a non-myopic loss, including multiple tokens ahead for language
modeling, improves the effect of thinking.
6. On multiple tasks, we demonstrate that thinking allows the LM to predict difficult tokens better than one trained on the same web text, improving with longer thoughts.
2
10 20 30 40 50 60 70 80 90 100
Training Step
6%
7%
8%
9%
10%
11%
GSM8K Accuracy (%)
# Thought, Ahead Tokens
24, 12
16, 8
12, 4
10, 4
8, 4
Baseline
(a) GSM8K
10 20 30 40 50 60 70 80 90 100
Training Step
34%
36%
38%
40%
42%
44%
46%
CommonsenseQA Accuracy (%)
# Thought, Ahead Tokens
24, 12
16, 8
12, 4
10, 4
8, 4
Baseline
(b) CommonsenseQA
Figure 2: Generalization Results. We evaluate the extent to which the model trained
with Quiet-STaR generalizes to directly answering problems that require reasoning. The
left plot (a) shows the zero-shot accuracy on GSM8K, while the right plot (b) shows the
zero-shot accuracy on CommonsenseQA, without any fine-tuning. In both plots, the x-axis
represents training steps, and each line corresponds to a different number of thinking tokens
used during Quiet-STaR training. The y-axis measures the zero-shot direct accuracy on the
respective datasets. We also include an inference normalized version of this plot in Figure 6.
2 Related Work
2.1 Reasoning in Language Models
There have been many works on training and exploiting language models to solve difficult
tasks by first training them to reason through them. For example, Rajani et al. (2019)
demonstrated that a pre-trained language model fine-tuned to output on human reasoning
traces before answering multiple-choice commonsense reasoning questions outperformed
one trained directly on answers. Shwartz et al. (2020) demonstrated that language models,
when provided with some scaffolding, can generate these helpful chain-of-thought solutions
without additional supervision. Later, Nye et al. (2021) demonstrated that “scratchpads”
required less scaffolding when the language models were more capable, a result later
reinforced by Wei et al. (2022b), emphasizing informal tasks, and further strengthened by
Kojima et al. (2022), demonstrating this behavior could be accomplished zero-shot. Most
recently, Wang & Zhou (2024) showed further that for commonsense-question answering,
one could force a language model to leverage chain-of-thought reasoning by preventing it
from emitting any valid answer tokens unless it was confident. However, once again, these
approaches only work for a question-answer dataset, and Wang & Zhou (2024) relies on
heuristics to identify when the model has output answer tokens. Somewhat like TRICE
(Phan et al., 2023), we use the relative improvements in the log-likelihood of the target text
across rationales as an estimate of quality, but we simply subtract the mean reward and do
not incorporate more complex control variates.
2.2 Training Language Models to Reason
One direction that researchers have used to train language models to reason or improve their
reasoning is training the language model on mined reasoning traces or reasoning-like data
(Rajani et al., 2019; Wei et al., 2021a; Lewkowycz et al., 2022; Chung et al., 2022; Gunasekar
et al., 2023). Although this approach has been demonstrated to be effective, it comes with
drawbacks. It requires either manual annotation, which is sensitive to the capability of the
annotators and is off-policy for the language model (i.e., the distribution of reasoning is not
text that the language model would otherwise likely have generated). This approach is also
expensive, difficult to scale, and provides no clear path to solving problems harder than
those that the annotators are capable of solving.
Another direction for teaching reasoning relies on a language model’s own generated
reasoning, which can be seen as building on a large body of literature on self-play (Silver
et al., 2017; Anthony et al., 2017; Polu & Sutskever, 2020). These include methods such as the
3
Algorithm 1: Quiet Self-Taught Reasoner (Quiet-STaR)
Input: Language model θ0, training steps num steps, sequence length l, thought length
t, learning rate α, batch size b, number of thoughts nthoughts, number of ground
truth tokens used for supervising each thought ntrue
Output: Language model θ that generates rationales to predict future text
for i = 0 to num steps do
Sample batch of sequences X of length l
h
init ← hidden statesθi
(X)
for j = 1 to l in parallel using attention mask do
log p
init
j:j+ntrue ← lm headθi
(h
init
j:j+ntrue
) // Predict next tokens
Tj ← generate tokensθi
([X:j
; <start thought>], t, nthoughts) // Generate thought
Tj ← [Tj
; <end thought>]
h
thought
j:j+ntrue ← hidden statesθi
([X:j
; Tj
; Xj:j+ntrue−1
])
log p
thought
j:j+ntrue ← lm headθi
(h
thought
j:j+ntrue
) // Predict next tokens w/ thought
wj:j+ntrue ← mixing headθi
(h
thought
j:j+ntrue
, h
init
j:j+ntrue
)
log p
talk
j ← wj:j+ntrue · log p
init
j:j+ntrue
+ (1 − wj:j+ntrue ) · log p
thought
j:j+ntrue
// Mix logits
L
NLL
j ← − log p
talk
j:j+ntrue
(Xj+1:j+ntrue+1)
rj = log p
talk
j:j+ntrue
(Xj+1:j+ntrue+1) − log p
talk
j:j+ntrue
(Xj+1:j+ntrue+1)
∇θL
REINFORCE
j ← −rj1[rj > 0] · ∇θ
log pθi
(Tj
|[X:j
; <start thought>])
∇θLj ← ∇θL
NLL
j + ∇θL
REINFORCE
j
θi+1 ← θi − α ∑
l
j=1 ∇θLj // Update model parameters
return θnum steps
Self-Taught Reasoner (Zelikman et al., 2022), which demonstrated that a language model
iteratively trained on its reasoning that led to correct answers could solve increasingly
difficult problems. Later work aimed to leverage additional information or assumptions
such as Huang et al. (2022) which demonstrated that the algorithm proposed in STaR could
still work if one assumed that the majority-vote answer was correct (although this has
a lower ultimate performance). Further work has generalized the results of Zelikman
et al. (2022), such as Uesato et al. (2022) which demonstrated additional usefulness to
“process-based” supervision where incorrect reasoning traces were filtered, recently VSTaR (Hosseini et al., 2024) that demonstrates that training a verifier to guide generation
also improves performance, as well as TRICE (Hoffman et al., 2024) which maximizes
the marginal likelihood of the correct answer given several reasoning traces per problem.
Finally, related work has also explored learning intermediate reasoning in the constrained
setting of making mathematical statements, where statements in the model’s intermediate
reasoning could be constrained to only be valid mathematical statements (Poesia et al., 2023).
We include further discussion of related reasoning works in Appendix F.
2.3 Meta-tokens
Recently, a growing body of work has demonstrated the usefulness of custom tokens
optimized to perform specific functions in the context of a neural network – for this reason,
they have also been referred to as “function vectors.” (Todd et al., 2023). One of the original
instantiations of this was prompt-tuning (Lester et al., 2021) (and relatedly prefix-tuning (Li
& Liang, 2021)), where the embeddings corresponding to the tokens of a prompt could be
optimized to better accomplish a task. Others have applied meta-tokens to compress long
prompts (Li et al., 2023; Jung & Kim, 2023) for efficiency. Most relevant to this work, Mu et al.
(2024) optimized a token such that, when the tokens after it could not attend to the tokens
before it (i.e., a context compression token), it would provide sufficient information to future
tokens. Although we do not focus on compression, we share the problem of learning a
4
token that affects attention and controls complex downstream behavior. In one related work,
Goyal et al. (2023) show that learning a single ”pause” token (essentially representing each
token as two tokens) improves LM performance. However, unlike the thought tokens in our
work, this pause token does not initialize a thought – instead, it can be seen as acting as the
entirety of the thought. We find that reasoning in language is significantly more helpful.
3 Problem Statement
In this work, we introduce an auxiliary ‘rationale’ variable between each pair of observed
tokens of the sequence. We then aim to optimize a language model with parameters θ with
the capacity to generate intermediate thoughts (or rationales) such that
θ
∗ = arg maxθ
Ex [logpθ (xi:n|x0:i
, rationaleθ (x0:i))]
Note that, in principle, this provides no advantage over an optimal language model that
already correctly models the language’s distribution over strings. Yet, in practice, extensive
prior work has shown that language models benefit from intermediate rationales on reasoning tasks (Nye et al., 2021; Zelikman et al., 2022; Wei et al., 2022b). Some work has aimed
to explain the effects of chain-of-thought reasoning, namely attributing it to “locality of
experience” (Prystawski et al., 2024). More broadly, reasoning allows a model to decompose
a challenging computation into smaller steps. In effect, we train the model to learn which
decomposition and planning steps are effective in predicting future text. Also note that we
formulate the objective as accurately predicting the remaining sequence, rather than only
the next token. Once again, for an optimal LM these would be equivalent. However we find
that the non-myopic formulation leads to a more effective loss for learning rationales.
4 Quiet-STaR
4.1 Overview
Quiet-STaR operates with three main steps (Figure 1):
1. Parallel rationale generation (think, Subsection 4.2): In parallel across n tokens
xi
in an input sequence x0:n, we generate r rationales of length t: ci = (ci1
, . . . , cit),
resulting in n × r rationale candidates. We insert learned 
 and

 tokens to mark each rationale’s start and end.
2. Mixing post-rationale and base predictions (talk, Subsection 4.3): From the hidden
state output after each rationale, we train a ”mixing head” – a shallow MLP producing a weight determining how much the post-rationale next-token predicted logits
should be incorporated compared to the base language model predicted logits. This
approach eases distribution shift early in finetuning, due to introducing rationales.
3. Optimizing rationale generation (learn, Subsection 4.4): We optimize the rationale
generation parameters (start/end tokens and LM weights) to increase the likelihood
of rationales that make future text more probable. We use REINFORCE to provide
a learning signal to rationales based on their impact on future-token prediction.
To reduce variance, we apply a teacher-forcing trick to include in the loss the
likelihood of predicting not only the token after the thought but also later tokens.
4.2 Parallel Generation
A key challenge in Quiet-STaR is efficiently generating rationales at each token position
in the input sequence. Naively, this would require a separate forward pass for each token,
which becomes computationally intractable for long sequences.
We allow for highly parallel generation by first observing that an inference pass of a language
model produces a probability distribution over the next tokens for all input tokens. Naturally,
this allows us to sample one next token from each token in the input. If one has generated a
successor from each token, it is not possible to simply continue with the original sequence.
For example, imagine predicting the next token after each token of “< bos > the cat sat” one
might generate “yes orange saw down” – each successor by itself is a reasonable next token
5
a’ b’ c’ d’
a’’ b’’ c’’ d’’
Base a b c d
Text
Thought
Token 1
Thought
Token 2
a b c d a’ b’ c’ d’
a
b
c
d
a’
b’
c’
d’
Base Text Thought Token 1
Generation Paths Parallel Inference Mask
Figure 3: Parallel Generation. By constructing an attention mask that allows all thought
tokens to pay attention to themselves, all preceding thought tokens within the same thought,
and the preceding text, we can generate continuations of all of the thoughts in parallel. Each
inference call is used to generate one additional thought token for all text tokens.
to a prefix of the sequence, but the list of tokens is a set of “counterfactual” continuations of
these prefixes. We can, however, leverage these continuations to generate hidden thoughts
for each observed token.
To do this efficiently, we cache each forward pass and concatenate a diagonal attention mask
to the previous attention mask: each generated token now attends to all of the tokens that
were used to generate it, as well as to itself (but not to token on other “counterfactual” paths).
Moreover, this parallelized next-sampling token procedure can be repeated arbitrarily many
times (or at least, until one runs out of memory). We visualize this procedure in Figure 3
and highlight additional ways to make this algorithm faster in Appendix B.
4.3 “Mixing” (Residual) Heads
When starting with a pre-trained model, thoughts will initially be out of distribution, and
hence harm language modeling performance. To smooth the transition to thinking, we
introduce a learned interpolation between the LM predictions with and without thoughts.
Given the end-of-thought token’s hidden state and the hidden state of the original text
token, the mixing head outputs a weight that determines the extent to which the postthought prediction logits will be used. We use a shallow multi-layer perceptron for this
head, outputting a scalar for each token. We include implementation details in Appendix A.
4.4 Optimizing Rationale Generation
4.4.1 Optimizing Start-of-Thought and End-of-Thought Tokens
The 
 and 
 tokens serve as learned meta-tokens that
control the model’s rationale generation. Optimizing the representation of these tokens,
especially the 
 token, is crucial but challenging due to the discrete nature
of the rationale tokens. We initialize the start and end token embeddings to the embedding
corresponding to the em dash, ”−−−”, which often appears in text data to denote a pause or
thought. This leverages the language model’s preexisting knowledge. In addition, to allow
these embeddings to be optimized more quickly, we apply a (hyperparameter) weight to
the gradients of these embeddings during the update step. Intuitively, the start thought
tokens can be understood as putting the model into a “thinking mode” and the end thought
token can be understood as telling the model when it’s done thinking.
4.4.2 Non-myopic Scoring and Teacher-forcing
Because we do not expect thoughts to be useful in predicting every token, we would prefer
the model’s reward to depend less on the exact next word in the text following the thought
and more on the following semantic content. There are two primary challenges here. First,
unlike in typical language modeling with transformers, only the thoughts corresponding to
6
d e f g h
<t> Thought </t> g h
Figure 4: Forward Pass and Teacher Forcing. We visualize a single forward pass of our
algorithm. Solid lines denote language model computation, while dashed lines indicate
tokens are inserted via teacher forcing, and the mixer represents the mixing head. In
particular, we visualize predicting three tokens ahead. Thought generation is shown in
more detail in Figure 1 and Figure 3.
a given next-token prediction receive a gradient from that prediction—a consequence of our
parallel sampling strategy. We could address this by adding loss terms for future tokens by
sampling the tokens before. However this would result in much higher entropy for language
modeling in general and lower-quality generated text, because it would train the LM to partially disregard its preceding tokens. Instead, we use the parallel attention mask to compute
the log probabilities of the true next tokens, applying teacher forcing by assuming the model
selected the correct next ground-truth token (as implicit in normal language modeling with
transformers). Note that the loss for each future token also depends on a mixing weight
computed from the end thought token and the previous observed token. The number of
future tokens included in the loss is a hyper-parameter. We apply the same teacher-forcing
technique to insert the start and end tokens. We visualize this procedure in Figure 4.
4.4.3 Objective
We use REINFORCE to optimize the likelihoods of the rationales based on their usefullness:
the log-likelihood of the ntrue true next tokens Xj+1:j+ntrue+1 under the language model given
previous observed tokens and a particular rationale (p
talk
j:j+ntrue
as shorthand for the mixed
prediction probabilities after thinking, see Algorithm 1). To reduce variance, we generate
multiple rationale continuations for each token in the input sequence (loosely inspired by
TRICE, Phan et al. (2023)). We thus define the reward rj
for each rationale Tj as the difference
between p
talk
j:j+ntrue
and the average across rationales for that token (p
talk
j:j+ntrue
):
rj = log p
talk
j:j+ntrue
(Xj+1:j+ntrue+1) − log p
talk
j:j+ntrue
(Xj+1:j+ntrue+1)
We then use this reward in a REINFORCE loss term to update the language model parameters θ to increase the likelihood of rationales that perform better than the average:
∇θL
REINFORCE
j = −rj
· ∇θ
log pθ (Tj
|[X:j
; 
])
We found it useful to exclude the negative reward from the REINFORCE loss term, as it led
to more stable training, though it may introduce some bias.
This loss term encourages the model to generate rationales that improve its predictions of
future tokens compared to the average prediction across all generated rationales for that
token. The gradients from this loss are used to update both the LM parameters and the
start-of-thought and end-of-thought token embeddings, with a (hyperparameter) weight
applied to the gradients of the start-of-thought and end-of-thought token embeddings to
accelerate their optimization. By iteratively optimizing these parameters, Quiet-STaR trains
the model to generate more useful rationales throughout training. Lastly, we also include a
log-likelihood loss, L
NLL
j
, to ensure that the LM learns to optimize the talking heads and
also receives a next-token prediction signal for the base LM head1
.
1Due to our linear mixing, equivalent to shifting the mixing weight toward the base prediction.
7
5 Experiments and Results
Intuitively, not all tokens require equal amounts of thought. For example, consider the
sentence “the person is run-”: although there is inevitably some probability of the token
being something other than “ing”2
, as a standalone sentence without context, additional
thinking is unlikely to improve a well-trained model’s prediction. Indeed, we conjecture
that for most chunks of most online text, additional thought has little to no impact. Indeed,
early in our exploration we observed that Quiet-STaR does not benefit all tokens equally.
Thus, we design our experiments to investigate whether our approach is useful in predicting
tokens that do require thought. We evaluate 1) whether Quiet-STaR improves a language
model’s ability to directly predict answers in datasets that require reasoning; and, 2) the
distribution of impacts resulting from thinking tokens. We conduct all of our experiments
starting with the base version of Mistral 7B (Jiang et al., 2023).
We perform most of our experiments by training on OpenWebMath (Paster et al., 2023),
a crawl that emphasizes more technical webpages. We selected OpenWebMath because
we anticipated that it would have a higher density of tokens that benefit from reasoning,
which our experiments support. We also evaluate Quiet-STaR on C4 (Raffel et al., 2020),
a widely used LM pretraining corpus with more diverse text, and again show significant
albeit smaller benefits.
5.1 Downstream Performance
In this subsection, we evaluate the extent to which Quiet-STaR improves the zero-shot
reasoning capabilities of the language model on CommonsenseQA (Talmor et al., 2018)
and GSM8K (Cobbe et al., 2021). On CommonsenseQA, we find that Quiet-STaR improves
performance by 10.9% compared to the base language model. As shown in Figure 2,
this improvement consistently increases with the number of tokens used in the model’s
rationales, indicating that more thorough reasoning through the thought tokens is translating
to better direct question-answering performance. Similarly, on GSM8K, Quiet-STaR results
in a 5.0% boost over the base model, and once again, performance scales with the length of
the rationales generated during Quiet-STaR training. For reference, in Figure 2, we include
a baseline corresponding to training the same model on the same dataset without thought
tokens. We observe that in multiple curves, performance appears to eventually deteriorate
– we anticipate that this is because we are not training on these downstream tasks, so the
roles of the thought tokens may change over time. We also find a benefit of our non-myopic
objective, which we discuss in Appendix D.
We find that training with Quiet-STaR on C4 (Raffel et al., 2020) also improves performance
on GSM8K (5.9% → 8.1%) and CommonsenseQA (36.3% → 42.6%) but by a smaller margin.
Specifically, for our C4 evaluation, we train Mistral 7B with 16 thought tokens and 4 true
tokens ahead and otherwise the same setup.
We can compare these improvements to those offered by pause tokens (Goyal et al., 2023),
which can be seen as a constrained version of Quiet-STaR where each token is represented by
two tokens and the second ”pause” token acts as the entirety of the thought. In particular, our
setup is most comparable to their pause token fine-tuning, as we also finetune a pretrained
model. Their results indicate that pause token fine-tuning also provides minor gains over
the base model on CommonsenseQA, they observed an improvement from 26.9% to 28.8%;
on GSM8K, Goyal et al. (2023) found that pause token fine-tuning harms performance.
Moreover, on both tasks (and the majority of their evaluated tasks), they observed that
additional thought tokens harmed performance. Moreover, they discuss the “lukewarm
effect of pause-finetuning a standard-pretrained model” (Goyal et al., 2023). This suggests
that allowing the model to generate multi-token rationales leads to more effective reasoning
compared to the single-token ”pauses”. Note however, that unlike Goyal et al. (2023), we do
not fine-tune on the downstream tasks.
Overall, these downstream results validate that training a language model to predict the
subtext between the lines of general text data can substantially improve its reasoning
2For example, in this very text, the token following “run” is “-”
8
1 2 3 4 5 6 7 8
Samples for Majority Vote
30
40
50
GSM8K Accuracy (%)
Accuracy on GSM8K vs Number of Majority Vote Samples
Quiet-STaR + CoT
Baseline (CoT)
Figure 5: Zero-shot performance on Quiet-STaR applied to chain-of-thought on GSM8K.
We visualize how using a Quiet-STaR trained Mistral model can improve chain-of-thought
performance. We use an 8-thought-token-trained model and use its internal thoughts to
improve the tokens in a zero-shot chain-of-thought (Kojima et al., 2022)
capabilities, even on datasets it was not explicitly trained on. The fact that longer rationales
consistently lead to better outcomes, and that Quiet-STaR outperforms the constrained
pause token approach, supports the notion that Quiet-STaR is successfully teaching the
model to leverage its own generated thoughts to reason more thoroughly about the input.
5.2 Improvement Distribution
As visualized in Appendix Figure 7, we find that on average there is little improvement
in the LM’s ability to predict arbitrary tokens. But, when we visualize the distribution of
relative improvements, there is a disproportionate improvement on more difficult tokens.
This reflects the idea that some text tokens are substantially harder and benefit more from
careful thought.
In Appendix Figure 8, we aim to provide some insight into the kinds of tokens where
the improvements occur. Namely, while thinking appears to help for many tokens in the
example, inspection suggests it disproportionately help to predict tokens where recalling
relevant information is useful, such as the name of an applicable theorem or the start of
the next step in a proof. Notably, this would align well with the framing proposed by
Prystawski et al. (2024).
5.3 Quiet-STaR and Chain-of-Thought
While there are natural parallels between chain-of-thought prompting and our approach,
they are orthogonal and complementary. In zero-shot chain-of-thought, a user actively
prompts the model to think ‘out loud’, otherwise using its ordinary production distribution
(Kojima et al., 2022); Quiet-STaR instead allows a model to think quietly at every token, with
a distribution trained to be useful. We investigate using silent, Quiet-STaR, rationales while
generating explicit CoT reasoning. Because our goal is generalist reasoning that requires no
task-specific input at all, we used a zero-shot prompt (“Let’s think step by step.”) without
in-context examples. Our experiments indicate that internal rationales allow the model
to generate more structured and coherent chains of thought, shown in Appendix E and
visualized in Figure 5. The majority vote accuracy over 8 samples (cot-maj@8) increases
from 40.6% to 47.7% with Quiet-STaR, as evaluated on a sample of 128 GSM8K test items.
Note that each chain-of-thought solution is sampled with temperature 0.7.
5.4 Examples
While there is no explicit regularization in Quiet-STaR for thoughts to be humaninterpretable, they are generated from the same transformer trained to model language,
hence likely to be at least partially understandable. We discuss why this design choice
benefits the training stability in Appendix I. For reference, we include examples of thoughts
generated that were helpful to the model in predicting future tokens in OpenWebMath. First,
in one case, recalling that one should start with magnesium to produce magnesium nitride
allows it to better predict that the first step of the procedure involves heating magnesium.
9
'
 # Magnesium reacts with nitrogen to form magnesium nitride. The chemical
formula for this reaction is Mg+N_2-> MgN_2. What is the product, or what
are the products, of this reaction?\n\nJan 12, 2016\n\nThe formula for
magnesium nitride is $M {g}_{3} {N}_{2}$.\n\n#### Explanation:\n\nAs do
many active metals, magnesium nitride can be
 1 --, so the
equation of the reaction that forms magnesium nitride is\n\n$Mg + N_2 \\to

 formed by heating the metal (fier'
In some cases, the most useful thoughts appear to be near-continuations that correspond
more closely to the target text, e.g.,
An integer $n$ is odd if $n = 2k+1$ for some integer $k$.\n\nTo prove that $A =
B$, we must show that $A \\subseteq B$ and $B \\subseteq A$. The first of
these tends to
 in some sense - to be the more difficult<|
endthought|> trickiest for students
Lastly, we include an example from answering CommonsenseQA. Notably, this thought
occurs while reading the question and hence was not used to predict the final answer.
'
 Q: Talking to the same person about the same thing over and over again is

\n\n(a) a one-to-one correlation\n\n(b) a one-to<|
endofthought|> something someone can what?'
6 Limitations
This work proposes a new framework for learning to reason, and in doing so explores
solutions to a variety of meta-learning challenges. However, to solve these challenges,
certain simplifications were necessary. For example, it would be valuable to understand
whether these techniques work when a model is trained from scratch. We have also only
applied Quiet-STaR to a 7 billion parameter model, albeit a powerful one. The same
techniques applied to a better model would likely yield disproportionately better results, as
has often been observed for gains from reasoning (Wei et al., 2022a).
Quiet-STaR results in a substantial overhead, generating many tokens before generating
every additional token. (See Appendix C for compute adjusted performance results.)
However, this can also be seen as an advantage: typically, a language model can generate the
next token based on the current context, and while there are techniques to improve sampling
quality, there is no general way to leverage additional compute to enhance next-token
prediction. In the current implementation we do not support dynamically predicting when
to generate, or end, a rationale. However, this would be a natural extension. For instance, if
the mixing head was a prediction from the base language model, before any thought, rather
than after the thought, one could apply a threshold to prevent generating thoughts that
would not be incorporated. We expect that this is a more difficult task, as predicting the
usefulness of a thought is simpler when one has already generated the thought.
7 Conclusion
Quiet-STaR represents a step towards language models that can learn to reason in a general
and scalable way. By training on the rich spectrum of reasoning tasks implicit in diverse
web text, rather than narrowly specializing for particular datasets, Quiet-STaR points the
way to more robust and adaptable language models. Our results demonstrate the promise
of this approach, with Quiet-STaR improving downstream reasoning performance while
generating qualitatively meaningful rationales. We believe this also opens many potential
future directions - for example, one may aim to ensemble thoughts in order to further
improve the predictions for future tokens. Moreover, if the language model can predict
when thought will be useful, for example by putting the mixing head before the prediction,
then the predicted mixing weight could be used to dynamically allocate compute during
generation. Future work can build on these insights to further close the gap between
language model and human-like reasoning capabilities.""