behind the scenes of paper writing

Pages of "Base Models Can Reason By Taking a Cue From Training Data".

I received many kind comments about the conceptual clarity and figure design of my most recent paper, "Base Models Can Reason By Taking a Cue From Training Data", so I wanted to share how this project came to be and the approach I took to writing it. One of my goals as a researcher is to help people understand our findings by communicating them clearly, and to help others in their own research by sharing the process behind those findings.

It is common to view communication as the final stage of research, as taught by the scientific method. For this paper, though, thinking about how to communicate our findings actually helped guide the research itself¹. Some may find it unconventional that I storyboarded the figures before we even had the experimental results needed to produce them. In fact, we designed our experiments around the figures, rather than the other way around.

This reflects what I value in a successful research project. In particular, I care about whether I can produce a change in understanding, first in myself, and then in other people.

#The conception of this paper

This project began during my summer as a research visitor at UC Berkeley, where I worked with Amil Dravid and Alyosha Efros. We were originally studying whether different reasoning capabilities corresponded to low-rank subspaces in weight space, and whether we could steer a base model’s behavior by moving within those subspaces. However, this direction remained unfruitful for the two months that we tried.

Instead, Alyosha (who believes that data is everything²) encouraged us to look at the language model's reasoning outputs themselves. He believed that we could distinguish correct from incorrect responses just from the tokens they contained. His longstanding intuition was that language models "reason" by using specific words and phrases to effectively index into relevant parts of their training data. At first, Amil and I didn't listen to his suggestion because we were pretty sure he was wrong. How could such complicated reasoning behavior be decoded just from the words a model used? Yet, Alyosha kept bugging us to try, and we eventually listened...

To our surprise, Alyosha was right! Looking at the base model’s responses, we found that those starting with certain tokens, such as “Okay” or “To,” were much more likely to be correct. We got an "Aha, I knew it!" from Alyosha³, and he was very happy to see evidence consistent with an intuition he had held for years. He later wrote in an online post:

Following Saussure and late Wittgenstein, I believe that commands like "reason step-by-step" don't actually need to carry any semantic meaning -- they could just be pointers into the relevant parts of the training data. Yet until now, I had no solid evidence.

We then tried forcing the base model to start with these particular tokens, while letting it generate the rest. Its reasoning performance improved substantially, which led us to hypothesize that post-training was making these effective openings more likely. We also noticed that many reasoning traces in the training data began with “Okay,” which made us interested in the data connection.

At this point, my mental imagery of the phenomenon began to form. Token cues could elicit reasoning in the base model, post-training might make those cues more likely, and their effect might come from associations learned during training. This was the picture I began storyboarding into figures before we ran any further experiments.

(The project later grew through an amazing collaboration with Rulin Shao, Kevin Farhat, and Sewon Min. This post focuses on its beginnings and my own approach to writing, but the finished paper was a collective effort!)

#Storyboarding through figures

Because I care deeply about scientific understanding, my measure of a good paper is whether the reader walks away thinking, “Oh, I understand something fundamental better now, and I can apply that understanding to a problem I care about!” I find this intensely gratifying. Storyboarding through figures helped me think concretely about how to get there: what would the reader need to see to understand the phenomenon, and in what order?

##Designing experiments around figures

Starting with figures helped me turn our intuitions into falsifiable hypotheses and design simple experiments to test them. For each sketch, I had to think about what I expected to happen, and sometimes the results did not match my expectations. This approach works particularly well for me because I tend to think visually. Below, I walk through some of my sketches alongside the final figures.

Sketch. Figure 2 was our starting point. We already had preliminary results showing that starting the Olmo-3-7B base model’s response with “Okay,” then letting it generate the rest, could substantially improve its accuracy on reasoning problems. This supplied opening is what we call a token cue. I sketched three conditions side by side, the base model without a cue, the base model with the “Okay” cue, and the reinforcement learning (RL) post-trained variant. I initially included both pass@1 and pass@16 to compare success with one attempt and with multiple attempts, inspired by work comparing base and RL models across different sampling budgets⁴.

Final figure. The message conveyed by this figure is whether a cue could bring the base model’s single-response accuracy close to its RL counterpart’s. Because each plot should convey a single idea, I kept pass@1 and removed pass@16, which answered a different question about success across multiple attempts.

Sketch. For Figure 3, I wondered whether RL improved accuracy partly by making effective openings such as “Okay” more likely. I sketched three side-by-side panels to test this and to visualize my hypotheses. First, would supplying the cue during RL help the model reach high accuracy with fewer training steps? Second, would the model choose that opening more often after RL? Third, when we compared the base and RL models’ probabilities for the next token along the same response, would the largest differences occur at the beginning? Together, these comparisons would examine what RL changed and whether supplying the cue could make training more efficient.

Final figure. My sketch led with training efficiency because we were initially interested in it as an application. As our focus shifted toward understanding the phenomenon, I wanted to lead with what RL changed in the model. The order of the panels should follow the explanation you want to communicate. We reordered the panels to move from where the models’ predictions differed, to which opening tokens became more likely, and finally to the implications for training efficiency. I also replaced the cue-probability bars with branching diagrams to show the opening-token sequences and their probabilities before and after RL.

Sketch. To understand why “Okay” worked in the first place, we turned to the training data. We had noticed that many reasoning traces in training began with this word. Thus, for Figure 5, we wanted to test whether an arbitrary word could acquire the same effect if it appeared in those contexts instead. We thought of an experiment where we would replace “Okay” with “Chicken” in the mid-training data, then resume training from the same checkpoint on the original and edited data. Thus, I sketched my hypothesis that the edit would make “Chicken” an effective reasoning cue across mid-training. I also expected “Okay” to lose its effectiveness.

Final figure. The results of this experiment only partly matched my sketch. Starting a response with “Chicken” now improved reasoning accuracy, but “Okay” still worked too. Thus, giving a new word an association with reasoning had not erased the old cue’s effect. We therefore tested a second edit that placed “Okay” at the beginning of questions in the training data to redirect its effect. After training with this data edit, the “Okay” cue was no longer effective at eliciting reasoning. The final figure includes both edits to separate creating a new cue’s effect from redirecting an existing one.

In the sketch, I originally showed how the cues’ effects emerged and diverged over training. Later, I realized that the main comparison was between the cues under each data edit, so I replaced the training curves with grouped bars. The axis of variation should make the comparison you care about clear. We also showed example edits and changes in next-token probabilities, connecting the edited data to the model’s opening tokens and reasoning performance.

Sketch. For Figure 8, we wanted to visualize how different cues might connect to different parts of the training data, without repeating an expensive counterfactual data-edit experiment for each one. I had a “flow” visual in mind, linking cues to different data sources, but had not yet worked out how to measure those connections. Sometimes, a sketch can start with a vague idea of what you want to understand and help you design an experiment around it.

Final figure. That visual gave us a concrete experiment to work toward. We compared the model’s internal representations while generating responses with its representations while processing training documents, looking for the most similar training examples for each cue. We also showed reasoning accuracy alongside these connections, helping readers see how different cues related to both different types of training data and different levels of performance.

##Making the figures

To make figures for this paper, I used Matplotlib. Below, I share some general tips about what to consider while creating figures for your paper⁵.

###Font

I like to choose my own font rather than relying on plotting defaults to give my figures a distinctive look. I use sans serif fonts for plot labels to distinguish them from the paper's text, and you can vary the font weight to emphasize important information. On the other hand, I use monospace fonts for code and literal text. For this paper, I used Graphik as the primary sans serif font and B612 Mono as the monospace font.

###Color

Regarding color, I chose Tailwind v4’s palette, which provides a range of hues with numbered shades from light to dark. It uses OKLCH, a color space designed to better reflect perceived differences in color. I recommend choosing a consistent color palette for your figures.

The Tailwind v4 color palette.

I kept the color coding consistent across figures. For example, I used neutral for the base model without a cue, amber for the reasoning cue, blue for RL, and red for the “Chicken” cue. This makes the same conditions easy to recognize across figures. I also used the same shade level for comparable elements, e.g., 400-level shades for bar fills and 600-level shades for their outlines, such as amber-400 with an amber-600 outline.

Overall, I like to use a wide range of hues in my figures to create visual contrast. For example, consider a variant of Figure 10 that uses a narrower color palette:

Although I did not use colormaps in this paper, they are worth considering when color represents a numerical quantity. For an ordered quantity such as probability, a sequential colormap progresses from low to high. A diverging colormap is useful when there is a meaningful midpoint, such as zero change relative to a baseline. In either case, consider how evenly the colors change perceptually, so the palette does not make small numerical differences look disproportionately large.

###Layout

The layout of empty space is often overlooked, but spacing can be used to deliberately create perceptual grouping. In the example below, the choice of spacing can make the same panels read as rows or columns, which emphasize different axes of variation that a reader would compare. For example, if each row shows one model under different conditions, smaller gaps within rows encourage readers to compare those conditions.

###Visual hierarchy

Visual hierarchy is about helping the viewer determine what deserves their attention. Size, color, and line weight can emphasize a result while keeping supporting information visible. To demonstrate this effect, consider the two plots below, which contain the same samples and fitted line. On the left, larger points draw attention to the samples and their departure from a linear trend, while on the right, a stronger line draws attention to the fit. Remember to carefully consider what your figures emphasize, and whether that is intended or not.

###Exporting

When exporting your figure, save it in a vector format such as PDF, with text preserved as text, rather than a raster format such as JPEG. This keeps lines and labels sharp when readers zoom in and lets them select and copy the text. If you later compress the paper’s PDF, choose settings that preserve vectors and fonts so the figures stay sharp instead of being downsampled as images.

##Teaser figure

I designed Figure 1 using Figma after we wrote the initial draft of our paper, which already contained our major results. Figure 1 is meant to encapsulate the phenomenon illuminated by our paper in one view. As I described earlier, my mental picture was: "Token cues could elicit reasoning in the base model, post-training might make those cues more likely, and their effect might come from associations learned during training." You can see my earliest version of the teaser figure with its final version side by side below.

Teaser figures are important because they illustrate what the reader should take away from the entire paper. (A good heuristic is, if the reader does not plan to read the paper, what do you want them to understand from the teaser figure?) However, they are also tricky to make because they are meant to distill an entire paper into one figure. Personally, I approach making teaser figures by closing my eyes and imagining how I would "see" the paper or how I would explain our findings through a chalkboard drawing.

In the case of our paper, I wanted to depict token cues as different branches the base model could take at the beginning of a response. Although next-token prediction can be viewed as an exponentially branching tree, I focused on the first few tokens because they can have a lasting effect on the behavior of the rest of the generation. I also wanted to connect these different behaviors to the training data. I initially considered using three separate panels for the base model, RL post-training, and the training data edit. However, I had a unified understanding of what RL and the data edit do to the base model, so I brought them together in one diagram. I pictured RL as making an effective reasoning branch more likely to be taken, while the data edit turned “Chicken,” an opening that was rarely taken and did not previously elicit reasoning, into another reasoning branch. I also decided to add the empirical results underneath the conceptual diagram to show how these changes translate into tangible effects on reasoning performance.

#Telling a story through words

Because there are many excellent guides about the various aspects of scientific writing⁶, I will keep this section brief and high-level.

We spent a ridiculous amount of time writing and rewriting this paper. One thing I learned while writing this paper is the importance of narrative—a form of storytelling. There should be two parallel narratives, in the figures and in the words. More importantly, each should be able to stand alone without the other.

Earlier, I described how I created a visual narrative. The purpose of the visual narrative is to let the results speak for themselves⁷, which is distinct from the goal of the written narrative. In an empirical research paper, the writing should help readers follow our reasoning about why the problem is interesting, why we chose these experiments, and what we conclude from the results. These three points may sound simple, but conveying them clearly can be surprisingly difficult. If your readers come away understanding all three, then you’ve done a great job.