# oladiy2504/where_do_gradient_methods_really_happen architecture

> Architecture diagram of the GitHub repository oladiy2504/where_do_gradient_methods_really_happen, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-07-31). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/oladiy2504/where_do_gradient_methods_really_happen
- Repository: https://github.com/oladiy2504/where_do_gradient_methods_really_happen

## Overview

#### Purpose and Entry Points
This is an ML optimization research repository for comparing optimizers and gradient/update projection methods across image, language, and transformer-style models. The primary experiment entry points are `experiments/run.py` and MLflow-integrated `experiments/mlflow_runner.py`; reusable execution logic is centralized in `src/experiments/runner.py`.

#### Experiment Orchestration
The experiment layer defines configurations, resolves tasks, runs seeded trials, and records histories/results. `experiments/configs.py` supplies experiment definitions, while `src/experiments/tasks.py` maps them to model/dataset workloads. `src/experiments/runner.py` is the core control boundary between configuration and training execution.

#### Model and Dataset Workloads
Model implementations cover MNIST MLPs, CIFAR CNN/ResNet/VGG models, GPT, and an SST transformer: `src/models/mlp_mnist.py`, `src/models/cnn_cifar.py`, `src/models/gpt.py`, `src/models/transformer_sst.py`. Dataset acquisition/loading and model-facing data preparation are concentrated in `src/models/data.py`; `experiments/prefetch_datasets.py` supports preloading datasets before runs.

#### Optimizer Implementations
The repository implements its optimizer variants rather than only wrapping framework defaults. Standard and experimental update rules reside in `src/optimizers/sgd.py`, `src/optimizers/adam.py`, `src/optimizers/muon.py`, and `src/optimizers/subzero.py`, with related algorithms such as AdamW, SGDM, MeZO, forward-gradient, and SWA in the same package. This layer consumes gradients and emits parameter updates, optionally exposing optimizer state such as momentum or Adam moving averages to projection logic.

#### Projection and Subspace Methods
Projection is a separate architectural layer applied to gradients, updates, or optimizer state. The interface/base abstraction is `src/projections/base.py`; implementations include eigenspace, Hessian, gradient-covariance, momentum-matrix, random, Stiefel, and two-sided projections in `src/projections/`. `src/projectors.py` likely provides composition or selection of these methods for the runner. This separation enables experiments such as dominant/bulk subspace treatment and tangent-space momentum projections without rewriting optimizers.

#### Results, Analysis, and Tracking
Runs persist experiment configuration, per-seed histories, validation histories, summaries, and method-specific result JSON under `data/`, exemplified by `data/momentum_cifar_adam_expavg_layerwise_tangent_energy080_switch_1000_end_3000_upd1_3seeds/`. Multi-seed aggregation is represented by `history_aggregate.csv` and `validation_aggregate.csv`. MLflow is the material external tracking service via `experiments/mlflow_runner.py`; plots and postprocessing are supported by `src/experiments/plotting.py` and notebooks under `notebooks/`.

#### Runtime and Reproducibility Boundary
The project is Python-based, with dependencies/environment defined by `pyproject.toml`, `environment.yml`, and `uv.lock`. Training likely relies on a tensor/autodiff runtime implied by the model and optimizer architecture, but the file tree does not identify a specific framework. Research outputs are retained in-repo as CSV/JSON data, notebooks, and LaTex report/poster artifacts in `report/` and `poster/`.

Main entry points: [`src/experiments/runner.py`](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/experiments/runner.py) (Experiment runner), [`src/experiments/tasks.py`](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/experiments/tasks.py) (Task resolver), [`src/models/data.py`](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/models/data.py) (Data preparation), [`experiments/run.py`](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/experiments/run.py) (CLI runner)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entry["Experiment Entry Points"]
  node_run["CLI runner<br/>Python entry point<br/>[run.py]"]
  node_mlflow_runner["MLflow runner<br/>Python entry point<br/>[mlflow_runner.py]"]
end

subgraph group_core["Execution Core"]
  node_configs["Experiment configs<br/>configuration definitions<br/>[configs.py]"]
  node_runner["Experiment runner<br/>orchestration<br/>[runner.py]"]
  node_tasks["Task resolver<br/>workload mapping<br/>[tasks.py]"]
end

subgraph group_workloads["Workloads"]
  node_data[("Data preparation<br/>dataset loader<br/>[data.py]")]
  node_prefetch["Dataset prefetch<br/>preload utility"]
  node_image_models["Image models<br/>[mlp_mnist.py]"]
  node_gpt["GPT model<br/>language model<br/>[gpt.py]"]
  node_sst_transformer["SST transformer<br/>transformer model<br/>[transformer_sst.py]"]
end

subgraph group_optimization["Optimization Methods"]
  node_optimizers["Optimizer variants<br/>update rules<br/>[sgd.py]"]
  node_projector_selector["Projector selection<br/>projection composition<br/>[projectors.py]"]
  node_projection_base["Projection interface<br/>abstraction<br/>[base.py]"]
  node_projection_methods["Subspace projections<br/>projection implementations"]
end

subgraph group_outputs["Tracking and Outputs"]
  node_results[("Run artifacts<br/>CSV / JSON results")]
  node_mlflow{{"MLflow<br/>external tracking service"}}
  node_plotting["Plotting and analysis<br/>postprocessing<br/>[plotting.py]"]
  node_notebooks["Research notebooks<br/>analysis notebooks"]
  node_reports["Reports and posters<br/>research artifacts"]
end

node_run -->|"loads"| node_configs
node_run -->|"starts"| node_runner
node_mlflow_runner -->|"starts"| node_runner
node_mlflow_runner -->|"tracks"| node_mlflow
node_configs -->|"defines trials"| node_runner
node_runner -->|"resolves"| node_tasks
node_tasks -->|"selects data"| node_data
node_tasks -->|"selects"| node_image_models
node_tasks -->|"selects"| node_gpt
node_tasks -->|"selects"| node_sst_transformer
node_prefetch -->|"preloads"| node_data
node_data -->|"feeds batches"| node_image_models
node_data -->|"feeds batches"| node_gpt
node_data -->|"feeds batches"| node_sst_transformer
node_runner -->|"executes training"| node_optimizers
node_runner -->|"configures"| node_projector_selector
node_projector_selector -->|"uses interface"| node_projection_base
node_projection_base -->|"implemented by"| node_projection_methods
node_projection_methods -.->|"projects gradients, updates, or state"| node_optimizers
node_runner -->|"persists histories and summaries"| node_results
node_results -->|"analyzed by"| node_plotting
node_plotting -->|"supports"| node_notebooks
node_notebooks -->|"informs"| node_reports

click node_run "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/experiments/run.py"
click node_mlflow_runner "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/experiments/mlflow_runner.py"
click node_configs "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/experiments/configs.py"
click node_runner "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/experiments/runner.py"
click node_tasks "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/experiments/tasks.py"
click node_data "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/models/data.py"
click node_prefetch "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/experiments/prefetch_datasets.py"
click node_image_models "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/models/mlp_mnist.py"
click node_gpt "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/models/gpt.py"
click node_sst_transformer "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/models/transformer_sst.py"
click node_optimizers "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/optimizers/sgd.py"
click node_projector_selector "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/projectors.py"
click node_projection_base "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/projections/base.py"
click node_plotting "https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/experiments/plotting.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_run,node_mlflow_runner toneBlue
class node_configs,node_runner,node_tasks toneAmber
class node_data,node_prefetch,node_image_models,node_gpt,node_sst_transformer toneMint
class node_optimizers,node_projector_selector,node_projection_base,node_projection_methods toneRose
class node_results,node_mlflow,node_plotting,node_notebooks,node_reports toneIndigo
```

## Components

### Experiment Entry Points

- [CLI runner](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/experiments/run.py): Python entry point · `experiments/run.py`
- [MLflow runner](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/experiments/mlflow_runner.py): Python entry point · Runs experiments with external MLflow tracking. · `experiments/mlflow_runner.py`

### Execution Core

- [Experiment configs](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/experiments/configs.py): configuration definitions · `experiments/configs.py`
- [Experiment runner](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/experiments/runner.py): orchestration · Resolves workloads, executes seeded trials, and records histories. · `src/experiments/runner.py`
- [Task resolver](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/experiments/tasks.py): workload mapping · `src/experiments/tasks.py`

### Workloads

- [Data preparation](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/models/data.py): dataset loader · `src/models/data.py`
- [Dataset prefetch](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/experiments/prefetch_datasets.py): preload utility · `experiments/prefetch_datasets.py`
- [Image models](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/models/mlp_mnist.py): MLP / CNN / ResNet / VGG · MNIST and CIFAR model implementations. · `src/models/mlp_mnist.py`
- [GPT model](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/models/gpt.py): language model · `src/models/gpt.py`
- [SST transformer](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/models/transformer_sst.py): transformer model · `src/models/transformer_sst.py`

### Optimization Methods

- [Optimizer variants](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/optimizers/sgd.py): update rules · SGD, Adam, Muon, SubZero, and related algorithms produce parameter updates. · `src/optimizers/sgd.py`
- [Projector selection](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/projectors.py): projection composition · `src/projectors.py`
- [Projection interface](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/projections/base.py): abstraction · Applies methods to gradients, updates, or optimizer state. · `src/projections/base.py`
- **Subspace projections**: projection implementations · Eigenspace, Hessian, gradient-covariance, momentum, random, Stiefel, and two-sided methods.

### Tracking and Outputs

- **Run artifacts**: CSV / JSON results · Stores configurations, per-seed histories, validation, summaries, and aggregates.
- **MLflow**: external tracking service
- [Plotting and analysis](https://github.com/oladiy2504/where_do_gradient_methods_really_happen/blob/main/src/experiments/plotting.py): postprocessing · `src/experiments/plotting.py`
- **Research notebooks**: analysis notebooks
- **Reports and posters**: research artifacts

## Connections

- CLI runner → Experiment configs: loads
- CLI runner → Experiment runner: starts
- MLflow runner → Experiment runner: starts
- MLflow runner → MLflow: tracks
- Experiment configs → Experiment runner: defines trials
- Experiment runner → Task resolver: resolves
- Task resolver → Data preparation: selects data
- Task resolver → Image models: selects
- Task resolver → GPT model: selects
- Task resolver → SST transformer: selects
- Dataset prefetch → Data preparation: preloads
- Data preparation → Image models: feeds batches
- Data preparation → GPT model: feeds batches
- Data preparation → SST transformer: feeds batches
- Experiment runner → Optimizer variants: executes training
- Experiment runner → Projector selection: configures
- Projector selection → Projection interface: uses interface
- Projection interface → Subspace projections: implemented by
- Subspace projections → Optimizer variants: projects gradients, updates, or state
- Experiment runner → Run artifacts: persists histories and summaries
- Run artifacts → Plotting and analysis: analyzed by
- Plotting and analysis → Research notebooks: supports
- Research notebooks → Reports and posters: informs