NexusPretraining

KO | EN

gitlite — search

#python #java #python3 #arduino #golang #machine-learning #rust #html #flask #javascript #seismology #nodejs

NexusPretraining

★ 17 Open GitHub ↗

The official code implement of <Nexus: Same Pretraining Loss, Better Downstream Generalization via Common Minima>

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Nexus: Same Pretraining Loss, Better Downstream Generalization via Closeness

[![arXiv](https://img.shields.io/badge/arXiv-2604.09258-b31b1b.svg)](https://arxiv.org/abs/2604.09258)

![Nexus Paper](https://raw.githubusercontent.com/huanranchen/NexusPretraining/main/experiments/figures/paper.png)

## Main Idea

We argue that optimizing "geometric closeness" or "gradient closeness" can achieve better downstream generalization than traditional pretraining methods, even when achieving the same pretraining loss.

The basic intuition is as follows:

```math
\underbrace{\mathcal{L}_{\mathcal{T}}(\boldsymbol{\theta}) - \mathcal{L}_{\mathcal{T}}(\boldsymbol{\theta}- \gamma \nabla \mathcal{L}_{\text{train}}(\boldsymbol{\theta})) }_{\text{decrease of downstream loss after one GD step on training set}} = \gamma \nabla \mathcal{L}_{\text{train}}(\boldsymbol{\theta})^\top \nabla \mathcal{L}_{\mathcal{T}}(\boldsymbol{\theta}) + \mathcal{O}(\gamma^2).
```

If one succeeds in optimizing the gradient similarity between the training set and the downstream set, then (in a first-order sense) the decrease in downstream loss after one gradient descent step on the training set is maximized. In other words, gradient similarity between the training set and the downstream set serves as a direct proxy for generalization.

One might argue that since we cannot access the downstream set during training, optimizing the closeness within the training set may not generalize to the downstream set. While true, we maintain that closeness within the training set is at least a necessary condition for closeness to the downstream set. If the gradients within the training set are dissimilar to each other, how can the overall training set gradient be similar to that of the downstream set? Therefore, we believe optimizing gradient closeness makes a substantial difference.

To optimize gradient closeness, we propose the Nexus Algorithm (or Nexus Approximator):

![Nexus Illustration](https://raw.githubusercontent.com/huanranchen/NexusPretraining/main/experiments/figures/cwa_illustration.png)

As illustrated in the figure, the Nexus update consists of the original update plus a Nexus regularizer, which corresponds to the gradient of the gradient similarity (in a first-order sense).

Furthermore, implementing Nexus requires only a few lines of code:

![Nexus Code](https://raw.githubusercontent.com/huanranchen/NexusPretraining/main/experiments/figures/nexus.png)

## Hyper-parameters

Compared to the baseline, we introduce only one additional hyper-parameter: the inner learning rate (`inner_lr`).

Note that the ratio of the loss to the regularizer is equivalent to `1 : O(inner_lr)`.

Therefore, our tuning strategy is to set the `inner_lr` as large as possible while maintaining a stable loss. In our experiments, we selected an `inner_lr` of 0.1.

Different datasets may require different `inner_lr` values, but we recommend adhering to the principle: "tune the `inner_lr` to be as large as possible while keeping the loss stable."

## Main Experiments

For 1B (520M excluding embeddings) and 3B (2.3B excluding embeddings) baselines, you can run the following commands:
```
NCCL_DEBUG=WARN accelerate launch --mixed_precision=bf16 --main_process_port 12300 train_all_baselines.py  --model_size 520m --optimizer Adam --eval_freq 1000
```
```
NCCL_DEBUG=WARN accelerate launch --mixed_precision=bf16 --main_process_port 12300 train_all_baselines.py  --model_size 2.3B --optimizer Adam --eval_freq 1000
```

For 1B and 3B Nexus models, you can run the following commands:
```
NCCL_DEBUG=WARN accelerate launch --mixed_precision=bf16 --main_process_port 12300 train_1b_adam_nexus.py  
```
```
NCCL_DEBUG=WARN accelerate launch --mixed_precision=bf16 --main_process_port 12300 train_3b_adam_nexus.py  
```

To evaluate the benchmarks, you can run the following commands:
```
python3 benchmark_companion_api.py  \ 
    --base_dir=/your_data_path/checkpoints/ \
    --project_name=huanran-Nexus \
    --run_name=520m-Cosine-Muon-1.0xCWA-CWABatch1-SGD-LR2e-3-InnerLR1e-3-SD (your wandb run name)
```

## Scaling Law
For 130M, 300M, 520M, 1.2B, and 2.3B baselines, you can run similar commands:
```
NCCL_DEBUG=WARN accelerate launch --mixed_precision=bf16 --main_process_port 12300 train_all_baselines.py  --model_size 130m --optimizer Adam --eval_freq 1000
```

```
NCCL_DEBUG=WARN accelerate launch --mixed_precision=bf16 --main_process_port 12300 train_nexus_scaling.py  --model_size 130m
```

## Limitations

Currently, applying the Nexus algorithm to Muon optimizers leads to optimization stagnation. The theoretical conditions under which Muon effectively operates remain an open question. We encourage the community to delve deeper into the convergence dynamics of Muon, which could ultimately help us address this compatibility issue.

NexusPretraining

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README