ZeitgeistLM

Inspiration

We really wanted to train a a model end-to-end to see if it was even possible during such a short period of time. We also thought trying to get a decoder-only model to "see" time was a worthwhile challenge after reading Calcifer Computing's TLM report.

What it does

We train ZeitgeistLM, a 124M-parameter decoder-only transformer from scratch on timestamped Reddit submissions and comments. ZeitgeistLM uses ordinary causal next-token prediction, augmented with a learned continuous-time prefix embedding projected from a scalar timestamp and prepended to the token sequence. This lets the model condition every prediction on time, interpolate smoothly between observed dates, and extrapolate beyond the training interval, all of which are capabilities that discrete learned time tokens do not naturally provide.

We provide an intuitive generation interface to let users see how the model has learned changes in memetic content on Reddit over time. We also use ZeitgeistLM to build an exploratory meme genealogy that uses hidden-state similarity plus chronology to build hypothetical lineage graphs between related posts.

How we built it

We use Reddit submissions and comments from fifteen meme-oriented communities[1] in the open-index/arctic corpus, with the data spanning from January 2011 to August 2022. We use 2011-2020 for training, 2021 for validation, and January-August 2022 for testing. We clean the data by removing deleted text, bot and moderation boilerplate, media placeholders, and URL-dominated entires. We also cap the length of comment threads to 100 posts and remove copypasta's to stop them from contaminating the training corpus too much.

We use GPT-2's BPE for tokenization, with 1025 token-windows whose date is the mean date of all its inputs. Since we use both submissions and comments to train, we pack data only of the same content type (submission/comment) and month into a single token window. If we have short incomplete tails, we drop them instead of mixing data across months.

Let (u) be a window’s UTC timestamp, with (u_0) at the beginning of 2011 and (u_1) at the end of 2020. Dates after 2020 yield (\tau > 1) and dates before 2011 yield (\tau < 0) which basically means that if you try to evaluate on this you extrapolate along the same scalar coordinate, which gives the extrapolation some credibility.

$$ \tau = \frac{u-u_0}{u_1-u_0} $$

We then just build a tiny MLP that converts this scalar into a 768-dimensional vector (the same dimensionality as a GPT embedding), so optimization magic can do its work. This vector occupies position 0 before the text tokens, which allows causal attention to attend to it without inserting a date string into the document.

# Temporal Vector MLP:
time_encoder = nn.Linear(1, config.n_embd)  # 1 → 768

# In the forward pass, for one normalized timestamp per example:
time_value = time_value.reshape(B, 1)       # [B, 1]
time_emb = self.transformer.time_encoder(time_value)[:, None, :]
                                                # [B, 1, 768]
x = torch.cat((time_emb, tok_emb), dim=1)

Everything else is just basically vanilla GPT stuff - 12 transformer blocks, 12 attention heads, a 1,024-token text context, tied input and output token embeddings, and a 50,304-entry vocabulary. Our continuous-time prefix embedding adds one transformation position, but loss is computed over the 1,024 text predictions. And again, the cool thing is that the model trains from scratch, which means that there's no anachronistic tendencies hidden in the non-existent base model that we then finetuned.

We sample packed windows uniformly across the manifest, and use gradient accumulation for effective batches of 262,144 tokens. Training uses AdamW, warmup followed by cosine learning-rate decay, gradient clipping at norm 1, and BF16 autocast on GPU. Checkpoints preserve model, optimizer, loader random states, and PyTorch random states for resumption.

The model run finished at 11,320 steps and 2.967B tokens processed, which is about one epoch of training over the cleaned dataset, sampled with replacement. Wall clock time was 3h 17m running on a single Modal H1000, including restarts and checkpointing. The final logged training loss was 2.2683, and the final validation estimate was 2.1927.

Challenges we ran into

The first model training run had its loss completely go haywire past 5,000 steps because we were updating the learning rate too aggressively. Should have eased off a bit! And we did. And then after that it was basically okay. But we underestimated the amount of time it would take to clean the data. I hate cleaning data.

Accomplishments that we're proud of

Zeitgeist is pretty funny sometimes! We also successfully modified the nanogpt architecture (even though it was a small modification), and actually saw a meaningful change in capabilities due to our temporal prefix embedding (ie extrapolation as free lunch that comes with it).

What we learned

Model training is not that hard, and a lot of architectural decisions are about ideas more than the technical implementation itself. Always keep an eye on your training loss because you never know when you might have to roll back to a previous checkpoint. And also there's a lot more to training a model than you might think, specifically CLEANING DATA OH MY GOD WHY DID THAT TAKE ALMOST AS MUCH TIME AS BUILDING THE MODEL I UNDERSTAND WHY DATA COMPANIES ARE VALUED SO HIGHLY NOW.

What's next for ZeitgeistLM

Better UI, more fun interactions to understand our tests, better written research draft! Would also like to SFT to turn it into a turn based model.

[^1]: r/memes, r/dankmemes, r/AdviceAnimals, r/me_irl, r/MemeEconomy, r/copypasta, r/dogelore, r/shitposting, r/okbuddyretard, r/2meirl4meirl, r/wholesomememes, r/comedyheaven, r/starterpacks, r/surrealmemes, and r/bonehurtingjuice