PolyPythias: Stability and Outliers across Fifty Language Model Pre-Training Runs

We thank the reviewer for their insightful comments and positive feedback. Below, we address the individual weaknesses (W) and questions (Q).

W1. “ $...$ model sizes studied in this work is overall small $...$ ”
We agree that models up to 410M parameters are small compared to many recent models; however, PolyPythias is still valuable for the research community. We repeat for convenience what we wrote in response to Reviewer yB47 (W1):

We agree that models up to 410M parameters are small compared to many recent ones. This choice reflects computational constraints, prioritising seed exploration and checkpoint granularity over model size. Despite this limitation, we believe that our suite is still a valuable resource for the research community (as we address in the introduction), as the range of sizes, seeds, and checkpoints allows researchers to study a variety of questions – just as the original Pythia suite continues to do. Those questions include questions about model stability w.r.t. initialisations and data order and, notably, about how similar the behaviour and internal states of smaller, more manageable models are to those of larger models (a key question in, e.g., Mechanistic Interpretability).

W2. “ $...$ predict outliers to make the training more stable $...$ ”
While we do not provide specific training guidance, PolyPythias is designed to help researchers explore factors that improve LMs. We agree that predicting outliers is an important question. In Sec. 5, we analyse factors behind unstable runs but find that predicting outliers from early checkpoints remains challenging and beyond this paper’s scope. To support further research, we provide pre-shuffled training tokens for studying data order effects $1$ and will release six additional seeds with decoupled data-ordering and weight-initialization seeds.

$1$ Palm: Scaling language modeling with pathways (Chowdhery et al., JMLR 2023)

Q1. “ $...$ how about comparing all the pairs of models trained with different seeds $...$ ”
Comparing all seed pairs would provide additional granularity. Still, we opted for the simpler approach of comparing against seed 0 to align with our goal of highlighting general trends in inter-seed agreement. Our experiments showed consistent trends across methods, whether using (simple) average agreement or Cohen’s kappa. While deeper analyses (e.g., pairwise comparisons, Fleiss’ kappa, or item-specific agreement $1$ ) could be valuable, we hope our work lays the groundwork for such explorations. We will include additional discussion of these considerations in the camera-ready version.

$1$ Learning Whom to Trust with MACE (Hovy et al., NAACL 2013)

Q2. “ $...$ summarizations at the end of each section $...$ ”
Thanks for the suggestion. We will update the manuscript accordingly.

We thank the reviewer for their insightful comments and feedback. We address the individual weaknesses (W), questions (Q), and suggestions (S) below.

We group W1, S1, and Q1 as they refer to the MDL setup.

W1. “ $...$ I am not convinced by the interpretation $...$ ”
In Sec. 4, we analyse how linguistic information is represented in the latent representations of the model across checkpoints and sizes. The probing framework (e.g., $5$ ) allows us to study this phenomenon by measuring how accurately linguistic information can be decoded from internal representations. Probing is operationalised as a classification task. Thus – as correctly pointed out – our stability results reflect probes' ability to solve certain tasks using models' representations. We will clarify this further in the camera-ready version.

S1. “ $...$ I would recommend to explain the MDL $...$ ”
Thanks for your suggestion. We justify the MDL setup below eq. 1 in the paper; however, we will revise the manuscript to clarify this point further. Other advantages of MDL are discussed in Q1.

Q1. “I find the sparsity inducing prior in equation 1) problematic $...$ ”
The prior is indeed kept constant across all setups: following prior work ( $1, 2$ ), we use a zero-mean Jeffreys prior that applies jointly to each probe, with a scaling factor that encourages sparsity $3, 4$ . We will add a detailed description in the appendix of the camera-ready version.

MDL allows one to distinguish between linguistic information being encoded vs. probe memorisation, thus solving the problem of random vectors being able to learn tasks (a problem that affects standard linear probes trained with cross-entropy; see $5-6$ ). Also, MDL is more robust to hyperparameter choices and enables consistent comparisons between different architectures ( $2$ ).

$1$ Information-Theoretic Probing with Minimum Description Length (Voita & Titov, EMNLP 2020)

$2$ Subspace Chronicles: How Linguistic Information Emerges, Shifts and Interacts during Language Model Training (Müller-Eberstein et al., EMNLP 2023).

$3$ Adaptive Sparseness using Jeffreys Prior (Figueiredo, NeurIPS 2001).

$4$ Bayesian Compression for Deep Learning (Louizos et al., NeurIPS 2017).

$5$ Information-Theoretic Probing for Linguistic Structure (Pimentel et al., ACL 2020).

$6$ Language Modeling Teaches You More than Translation Does (Zhang and Bowman, BlackboxNLP 2018).

$7$ Designing and Interpreting Probes with Control Tasks (Hewitt and Liang, EMNLP 2019). Max

W2. “ $...$ use gender bias data sets released for other tasks. $,,,$ ”
We chose to focus on the benchmarks that capture different manifestations of gender (bias) during training: CrowS-Pairs tests semantically more complex gender stereotypes (e.g., “men/women cannot drive” vs. “nurse→she, doctor→he”) while BLiMP and Simple Co-occurrence Bias rely on simple gendered pronoun resolutions. Expanding the analysis to datasets from other tasks (e.g., MT) might provide additional context but it further complicates interpretations as it is unclear whether benchmarks for different tasks measure the same constructs $1,2$ . Moreover, previous work has identified many potential sources of benchmark “unreliability” $3,4,5,6$ that require comprehensive analyses (e.g., $7$ ). While valuable, these lie outside this paper's scope, but we will expand our discussion of these issues in the camera-ready version.

$1$ Language (technology) is power: A critical survey of "bias" in NLP (Blodgett et al., ACL 2020)

$2$ Undesirable biases in NLP: Addressing challenges of measurement (van der Wal et al., JAIR 2024)

$3$ Efficient Benchmarking (of Language Models) (Perlitz et al., NAACL 2024)

$4$ The MultiBERTs: BERT Reproductions for Robustness Analysis (Sellam et al., ICLR 2022)

$5$ Underspecification presents challenges for credibility in modern machine learning (D’Amour et al., JMLR 2022)

$6$ PATCH--Psychometrics-AssisTed benCHmarking of Large Language Models: A Case Study of Mathematics Proficiency (Fang et al., ArXiv 2024)

$7$ fl-IRT-ing with Psychometrics to Improve NLP Bias Measurement (Bachmann et al., Minds and Machines 2024)

W3. “ $...$ motivate why you use feature maps $...$ ” + S2. “ $...$ If previous work explained the limitations $...$ ”
We will explain the individual features used to construct the training map at the beginning of Sec. 5.

A training map is useful because it allows one to aggregate information from a whole list of low-level features. One could heuristically do this, but training maps are more principled. We understand that the individual metrics (e.g., loss) can provide intuitive information about the training dynamics, but they do not provide a comprehensive picture (and may be a result of the aforementioned low-level features).

We thank the reviewer for their positive feedback. We have carefully addressed the comments and suggestions in the other reviews. If our responses to other reviewers alleviate any remaining concerns, we hope the reviewer will consider increasing their rating.

We thank the reviewer for their insightful comments and positive feedback. Below, we address the individual weaknesses (W) and questions (Q).

W1. “As noted in the paper, the size of the model trained by the authors is limited $...$ ”
We agree that models up to 410M parameters are small compared to many recent ones. This choice reflects computational constraints, prioritising seed exploration and checkpoint granularity over model size. Despite this limitation, we believe that our suite is still a valuable resource for the research community (as we address in the introduction), as the range of sizes, seeds, and checkpoints allows researchers to study a variety of questions – just as the original Pythia suite continues to do. Those questions include questions about model stability w.r.t. initialisations and data order and, notably, about how similar the behaviour and internal states of smaller, more manageable models are to those of larger models (a key question in, e.g., Mechanistic Interpretability).

**W2. “It is also still unclear if the phenomena $...$ ” + Q3. *“ $...$ observed state transisions in the HMM analysis are generic tendency or not $...$

You are correct: PolyPythias capture (by chance) training instabilities (i.e., “loss spikes”) for two of the 50 runs considered. In fact, we proactively tried to minimise unstable runs to obtain as many well-trained models as possible. Importantly, in creating the multi-seed collection, we want to allow researchers to study a broader definition of stability (beyond “loss spikes”). For example, how stable are the learned representations across model size and seed? Or, how consistent are the mechanisms/circuits that emerge during training? How reliable are existing benchmarks w.r.t. the data order/model seed? Answers to these questions are orthogonal to training instabilities as signified by “loss spikes”. Finally, we note that understanding the cause of “loss spikes” would require experiments specifically designed for this purpose (e.g., $1,2$ ), which are outside the scope of this paper. We will make this clearer in the next version.

$1$ Latent State Models of Training Dynamics (Hu et al., TMLR 2024)

$2$ Stabilizing transformer training by preventing attention entropy collapse (Zhai et al., PMLR, 2023)

Q1. “ $...$ it would be useful to plot all data points.”
To limit computational costs, for each size and seed, we evaluate performance on a subset of the available checkpoints–(log-spaced) steps 0,1,2,...,512, 1k, and from step 3k onwards we choose every 10k-th step up to 143k (Sec 3.1). For the 14M model, we ran initial probing experiments for all checkpoints. While not representative of all setups, the resulting trajectories were highly consistent with what we observed for the final subsampled checkpoints. We hope this clarifies the questions.

Q2. “ $...$ save and analyse checkpoints at more frequent intervals $...$ “
We do agree that more checkpoints would be useful. However, given the size of this release (i.e., 7000 checkpoints) we kept the same checkpointing frequency as the original Pythia suite. This rate is consistent or higher than prior efforts for publishing intermediate checkpoints with comprehensive coverage (e.g., MultiBERTs $1$ , BLOOM $2$ , Mistral $3$ ).

$1$ The MultiBERTs: BERT Reproductions for Robustness Analysis (Sellam et al., ICLR 2022)

$2$ BLOOM: A 176B-Parameter Open-Access Multilingual Language Model (Scao et al., BigScience Workshop 2022)

$3$ https://github.com/stanford-crfm/mistral (Mistral, 2022)

Q4. “ $...$ I'm curious if authors observed similar trends to that I noted here or only the two the paper mentioned.”
This is indeed an interesting question. In our analysis (Fig 4), we observe that runs that diverge try to recover by returning to a previous training state. Due to space constraints, we have not investigated this further in the paper. However, addressing this type of question is exactly the goal of PolyPythias. We will extend the discussion of this point in the camera-ready version appendix.