oladiy2504/where_do_gradient_methods_really_happen

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/.