HessianAnalysis

Experiments on whether independently initialized causal language models develop similar gradient, function-update, subspace, and curvature geometry during training.

What this repository does

The main experiment trains several small GPT-2- or Llama-style causal language models from different random initializations while keeping a fixed set of measurement batches. At checkpoints it records:

• training and held-out evaluation loss;
• cross-seed gradient alignment;
• cross-seed similarity of local function-space updates;
• gradient-covariance concentration and dominant-subspace overlap;
• finite-difference and sparse Hessian diagnostics;
• Hessian trace, empirical Fisher trace, and Jacobian Frobenius estimates; and
• optional local learning coefficient estimates through devinterp.

The fixed reference batches are intended to separate changes in optimization geometry from measurement noise caused by changing probes.

Research question and scope

The concrete question here is whether models trained on the same distribution but from different seeds converge toward similar local training geometry—and, if so, whether that similarity is more visible in parameter gradients, low-dimensional subspaces, function-space effects, or curvature.

This is an exploratory measurement framework, not evidence that neural networks form organisms, undergo natural selection, or become mesa-optimizers. Those ideas helped motivate an earlier line of thought, but this public repository deliberately contains only the operationalized experiment and its technical background.

Status

Research prototype. The training, plotting, and multi-run comparison command-line tools are functional, checkpoint/resume behavior has automated coverage, and the metric definitions are documented. There are not yet benchmark result sets, a paper, or a stable Python API. Several curvature measurements are computationally expensive; start with a debug-sized model and sparse measurement intervals.

Repository map

• gradient_alignment_experiment.py — trains multiple seeds and records checkpoint metrics and restart state.
• plot_gradient_alignment_results.py — renders the results of one run.
• compare_gradient_alignment_runs.py — compares metrics across run directories.
• tests/test_checkpoint_resume.py — verifies exact continuation, stream replay, and shared-batch behavior on tiny CPU models.
• metrics_reference.pdf — definitions and interpretations of the logged metrics (TeX source).
• hessian_fisher_background.pdf — technical background on Hessian and Fisher quantities (TeX source).

Run directories contain config.json, metrics.json, checkpoint-latest.pt, and generated plots. They are ignored by Git. The checkpoint is replaced atomically only after a complete measurement, so an interruption leaves the previous restart point intact.

Installation

Python 3.11 or newer is recommended. The pinned NumPy build in the validated cloud environment requires Python 3.11 or newer.

python3 -m venv .venv
source .venv/bin/activate
python3 -m pip install -r requirements.txt

For a reproducible cloud environment, use requirements.lock.txt or begin with a compatible PyTorch/CUDA image and install the remaining pinned packages from that file. Each run records the actual Python, package, CUDA, GPU, and Git versions in config.json.

The required packages are NumPy, Matplotlib, PyTorch, Hugging Face Datasets, and Transformers. Install devinterp separately only if you want sparse local learning coefficient estimates.

Dataset access and tokenizer/model downloads are handled by Hugging Face, so the first run requires network access. A CUDA GPU is strongly recommended for non-debug experiments. The bfloat16 mode requires compatible hardware; CPU runs should use float32.

Quick start

Start with the small GPT-2 preset, a modest number of seeds, and infrequent expensive measurements:

python3 gradient_alignment_experiment.py \
  --output-dir results_debug \
  --run-label c4-gpt2-debug \
  --model-name gpt2-debug \
  --num-seeds 4 \
  --train-steps 1000 \
  --checkpoint-every 250 \
  --hessian-eig-every 1000 \
  --batch-size 8 \
  --eval-batch-size 8 \
  --cov-num-batches 8 \
  --cov-batch-size 1 \
  --block-size 64 \
  --max-eval-blocks 128 \
  --dtype float32

The default dataset is English C4 with streaming enabled. Use --help for model presets, optimizer choices, dataset controls, metric frequencies, and architecture overrides.

RedPajama example

The experiment can infer common text fields such as text, raw_content, and content:

python3 gradient_alignment_experiment.py \
  --output-dir results_redpajama_sample \
  --run-label redpajama-gpt2-debug \
  --model-name gpt2-debug \
  --dataset-name togethercomputer/RedPajama-Data-1T-Sample \
  --dataset-config "" \
  --train-split train \
  --eval-split train \
  --shuffle-eval-stream \
  --num-seeds 4 \
  --train-steps 1000 \
  --checkpoint-every 250 \
  --hessian-eig-every 1000 \
  --dtype float32

For a dataset with another schema, pass --dataset-text-field FIELD.

Cloud execution and resume

--train-steps is required: a metered run cannot accidentally train forever. At every completed measurement, checkpoint-latest.pt captures all model and optimizer states, random-number-generator states, fixed reference tensors, temporal comparison vectors, accumulated metrics, and the training-stream position. Scientific arguments are recorded in a resume signature; changing one accidentally causes resume to fail rather than silently mixing runs.

On a RunPod-style machine, keep the repository, Hugging Face cache, and run directory on persistent storage such as /workspace, not the disposable container filesystem. A bounded pilot might use:

python3 gradient_alignment_experiment.py \
  --output-dir /workspace/hessian-runs/c4-gpt2-debug \
  --run-label c4-gpt2-debug \
  --model-name gpt2-debug \
  --num-seeds 4 \
  --train-steps 1000 \
  --checkpoint-every 100 \
  --hessian-proxy-every 100 \
  --hessian-eig-every 500 \
  --llc-every 0 \
  --dtype float32

Resume with the same scientific arguments and either the same or a larger finite step budget:

python3 gradient_alignment_experiment.py \
  --resume-from /workspace/hessian-runs/c4-gpt2-debug/checkpoint-latest.pt \
  --model-name gpt2-debug \
  --num-seeds 4 \
  --train-steps 2000 \
  --checkpoint-every 100 \
  --hessian-proxy-every 100 \
  --hessian-eig-every 500 \
  --llc-every 0 \
  --dtype float32

For Hugging Face streaming datasets, resume deterministically rebuilds the stream and replays the already consumed batches before continuing. That can be slow after a long run and still depends on the remote dataset snapshot staying available. A fixed --local-text-file is the stronger option when exact long-horizon restartability matters.

Run the cloud-readiness checks locally without a GPU or dataset download:

python3 -m unittest discover -s tests -v

Plotting and comparison

Re-render one run:

python3 plot_gradient_alignment_results.py --input-dir results_debug

Compare multiple runs using the labels saved in each config.json:

python3 compare_gradient_alignment_runs.py \
  --input-dirs results_c4_sgd results_redpajama_sample \
  --output-dir comparison_plots

Select particular metrics when the full dashboard is too crowded:

python3 compare_gradient_alignment_runs.py \
  --input-dirs results_c4_sgd results_redpajama_sample \
  --labels c4-sgd redpajama-sgd \
  --metrics eval_loss function_update_cross_batch_cosine \
    gradient_to_hessian_top1_alignment hessian_top1_eigenvalue \
  --output-dir comparison_plots_selected

Interpretation cautions

• Parameter-space alignment is sensitive to symmetries and reparameterizations; function-space measurements are included partly for that reason.
• A shared dominant subspace does not by itself establish a shared algorithm or causal mechanism.
• Sparse Hessian, trace, Jacobian, and local-learning-coefficient estimates depend on probe choices and numerical approximations.
• Cross-seed convergence should be compared against initialization, optimizer, dataset, architecture, and measurement controls.
• The repository currently provides instrumentation rather than a validated empirical conclusion.

License

Released under the MIT License.