General plan/list for things that need to be implemented for better experiment/analysis workflow

[adamimos]

The purpose of this issue is to have a conversation/brainstorm about the things that should be implemented to get Simplexity into a more working form for general use, with reproducibility, organization, and ease of use. Some things off the top of my head that come up:

• Seems like everything should be MLFlow, including artifacts, saved ckpts, even analysis results?
• @ealt mentioned a tagging system, in order to deal with determining what should be saved (or not)
• Currently we have to set seq_len, n_ctx, vocab size, bos, etc. and its hard to think through how they all relate. In the standard case defining one or two of those should set the rest.

[kylejray]

bos, seq_len, n_ctx, etc...

I propose standard case would be that you need to define just 2 things

1. seq_len
2. boolean for if BOS is used or not.

vocab_size, n_ctx, bos token all generated under the hood based on the above.

Disclaimer: I've just started using this stuff, so I might be accidentally leaving out some obvious and important use cases.

[adamimos]

@ealt Did you mention that you already worked on an mlflow persister? Or is that work to be done?

[loren-ac]

Some thoughts:

1. Yes to MLFlow being default for experiment/results logging -- including model params
2. I think worth investing some time to think about best ways to use tags but not sure I understand what you mean by using them to determine what should be saved or not. I currently use them for various metadata things including referencing other MLFlow runs / artifacts the run depends on [1], setting the "experiment type" (e.g. model stitching, model training, etc.), and I have played around with the idea of using them to implement the organizational notion of a "project" (which I don't think MLFlow handles natively in any opinionated way).
3. We talked once about codifying some notion of the reproducibility of a run that is saved to MLFlow. I still like this idea and think that as the number of experiments we run grows it will be useful to have. I'm not sure exactly what should be logged for something to be considered and ultimately flagged (perhaps using tags) as reproducible but would be interested in having a discussion about this.
4. I know we talked about something like an Experiment object in the past. I'm less sure of the utility at the moment especially if MLFlow becomes sort of a single source of truth for a given experiment but I could be convinced that a version of this would be useful. I'm imagining something that essentially takes a resolved hydra config (or an MLFlow run_id) and then has something like a run method and returns something like a results object. Maybe would be nice to have especially if you are messing around in a notebook and you want to make little tweaks to experiment config or mess around with other analysis/visualization stuff?

[1] e.g. a training run that produced the model that my stitching experiment is using -- there might be better ways of doing this / might be less relevant once we start saving model params to MLFlow

[loren-ac]

I also wonder how any of the above (especially reproducibility and logging stuff) will need to be adapted for when the data we are using in our experiments doesn't come from a ghmm implemented in python but from some text, video, etc. created by a human.

[hrbigelow]

Just a general comment - I think just put the word out to everyone that if they have some feature request or difficulty with simplexity, encourage them to open an issue directly on the github with a specific description of the problem. And link particular functions / source code lines in the simplexity repo relevant to the issue (if they exist)

This will serve as a nice hub for everyone to see how people are using it and what is missing.

[ealt]

Our process for producing research requires a couple tools:

• Code: GitHub, with separate library and experiment repos, different branches, etc.
• Data: MLFlow should be used to track all the experiment details and artifacts (including model checkpoints, the outputs of analysis, etc.).
• Publications: Overleaf for writing up reports on completed projects (plus maybe someday something equivalent to Anthropic's Transformer Circuits site 👀)
• Documentation:
  • description of Simplex's overarching research mission/goals
  • general areas of strategic focus
  • collections of specific research ideas
  • something like an electronic lab notebook for individual projects

Maybe beyond the intended scope of this issue, but I do think we would benefit from improving our approach to generating and organizing documentation of project ideas/works in progress

Experiments

• I think the highest level of organization is the experiment
• We should think carefully about the scope of each experiment (there is some risk of thinking every new piece of work is a progression on what came before and end up with a couple sprawling experiments exploring several distinct ideas)
• It would be nice to have consistent naming conventions for experiments
• We should also add experiment descriptions that are concise but clearly demarcate them from related experiments
• We can also tag experiments. Some tag proposals:
  • project_doc: url to some unified document which describes and links to everything relevant to the porject
  • research_pillar: factored_representations, beliefs_in_llms, interp_tools, universal_representations, rl, etc.
  • status: in_progress, complete, published, deprecated
  • scale: toy, llm

Runs

• Under the umbrella of each experiment is a series of runs which can be thought of as related to the execution of some piece of experimental code (though something like a hyperparameter search will generate multiple runs)
• This can include tracking the data associated with training a model, running an analysis, or producing a figure or similar artifact
• The goal is for an MLFlow run to contain all the information necessary to reproduce it
• Run descriptions can be helpful, ideally they would answer questions like:
  • How does the particular run fit into the context of the broader experiment?
  • What was the goal for the particular run?
  • What distinguishes it from other runs in the experiment?
  • What are the key takeaways from the run?
  • What are the steps to reproducing this run?
• Tags become increasingly necessary as we plan support running larger number of runs per experiment
  • research_step: hyperparam_search, training, sft, rlvr, eval, analysis, generating_pretty_figures
  • retention: temp, indef
  • use logger.log_git_info() for reproducibility
  • tags can be used to filter and group runs, so whatever filtering/grouping categories make sense for your project
• Metrics can be grouped by prefixing a group name + /, e.g. train/loss, train/accuracy, test/loss, test/accuracy, when logging lots of metrics making use of these groups makes navigating the UI much nicer
• Parameters - always log the config values as parameters
• Artifacts
  • always log associated resolved configs
  • artifacts can also be grouped hierarchically, make use of that especially when logging similar artifacts (e.g. some equivalent data ever n steps)

Models

• Should be logged to MLFlow
• I'm still learning about MFLow's support for integrating models into its ecosystem

[hrbigelow]

Just one note on reproducibility - if the is_dirty flag is true when doing logger.log_git_info(), it's not possible to restore code state from that information.

[loren-ac]

@ealt on first pass a lot of what you laid out really resonates with me. I want to spend some more time going through everything you said in more detail but I wanted to share one high level thought/question that Im curious for your perspective on. While experiment is the highest native level of organization in MLFlow, I feel like there is another conceptual level above that which might be called the "Project" and a project might require organizing information into multiple experiments which each might have multiple runs. Perhaps your notion of Experiment as laid out above is actually the same as my notion of Project but I would be curious to get aligned on this. Im finding myself wishing there were another level above Experiment (in the MLFlow sense) right now but maybe there is a way to reframe things such that it's not needed. Curious for yours (and others!) takes on this.

[ealt]

@ealt on first pass a lot of what you laid out really resonates with me. I want to spend some more time going through everything you said in more detail but I wanted to share one high level thought/question that Im curious for your perspective on. While experiment is the highest native level of organization in MLFlow, I feel like there is another conceptual level above that which might be called the "Project" and a project might require organizing information into multiple experiments which each might have multiple runs. Perhaps your notion of Experiment as laid out above is actually the same as my notion of Project but I would be curious to get aligned on this. Im finding myself wishing there were another level above Experiment (in the MLFlow sense) right now but maybe there is a way to reframe things such that it's not needed. Curious for yours (and others!) takes on this.

You're right, "experiment is the highest native level of organization in MLFlow" is a more precise phrasing but for our purposes it does make sense of have even higher level groupings. I think the "Research Pillars" outlined in Paul and Adam's Simplex research outline doc are probably the highest level, then below that is something like a project which probably has the scope of an academic publication, and a particular project can be potentially comprised of multiple experiments (often something like the scope of different sections within a paper). I think a helpful heuristic is if it would make sense to group or compare data from different runs and describe them in a single section of a paper they should be within the same experiment, and if such a direct comparison doesn't make sense then maybe they should be separated into different experiments. To organize these higher levels of grouping in the context of MLFlow we should use experiment tags.

[ealt]

Beyond MLFlow I think we need to think how best to use documents to organize our research. I am imagining something like a hierarchy of docs that contain the links to navigate to every relevant piece of context and data we produce. There could be:

• Organization level: Simplex research mission and goals doc which would contain or link to our high level research outline
• Pillar level: Research outline doc listing different pillars of focus with descriptions for each as well as links to docs for each project under the relevant pillar
• Project level: Canonical docs for each project that would then describe the project goals, hypotheses, proposed experiments and would link experiment docs. Eventually links to an overleaf doc when it is time to write up results, and finally link(s) to where the project is published.
• Experiment level: Canonical docs of each experiment would describe the experiment in context with the broader project, description of research methods (data, models, analysis, etc.), link to the associated research branch on Github, link to the associated experiment on Databricks. MLFlow experiments should also link back to this corresponding experiment doc via the description/tag. It could also function as a lab notebook documenting the iterative process of ideas to explore and how things turned out (potentially linking to individual MLFlow runs).

With this structure we should have all the information relevant to our research neatly organized in easy to navigate graph of linked pages (documents, MLFlow, Overleaf, arxiv, etc.)

Right now to the extent that we do document research ideas and progress at all it has been in isolated Google docs. A simple improvement would be a shared drive of docs following some structure such as what I proposed, but I think a better solution would be to use some other tool. I personally like Obsidian which we could use in conjunction with Github. Advantages of that approach include all the benefits of version control via git, the sharing and tooling built into Github, all the features of Obsidian, and documents that exist as a directory of plaintext files which are easy to expose to all the AI agents that can interact with our local docs.