简而言之： 我们通过使用增强型优化器放大参数梯度的低频成分，从而加速了 Grokking 现象。

摘要： 机器学习中一个令人困惑的现象被称为 Grokking，即在对训练数据几乎完美过拟合之后，经过数十倍的迭代才实现泛化。针对机器学习从业者关注的这种长时间延迟问题，我们的目标是加速处于 Grokking 现象下的模型的泛化过程。我们将训练过程中参数梯度随时间变化的一系列信号视为随机信号，并将其分解为两个成分：快速变化、导致过拟合的成分，以及缓慢变化、促进泛化的成分。基于这一分析，我们仅需几行代码即可放大梯度中的慢速变化成分，从而将 Grokking 现象加速超过 50 倍。实验表明，我们的算法适用于图像、语言和图等多种任务，使得这种突然泛化的奇特现象具有实际应用价值。

使用方法

安装

Grokfast 除了 PyTorch 外，无需其他额外依赖包。文件 requirements.txt 仅用于复现文章中的实验，具体说明见下文的 复现部分。

指导说明

只需在调用优化器之前插入一行代码，即可应用 Grokfast。

1. 从我们的仓库下载单个文件 grokfast.py。

wget https://raw.githubusercontent.com/ironjr/grokfast/main/grokfast.py

2. 导入辅助函数。

from grokfast import gradfilter_ma, gradfilter_ema

3. 在训练循环 之前 插入以下一行。

grads = None

4. 在 loss.backward() 和 optimizer.step() 之间，插入以下其中一行。请确保 model 是 nn.Module 类型，并且在训练循环开始前已正确初始化 grads：

# ... 在优化循环中。
loss.backwards() # 计算梯度。

### 选项 1：Grokfast（有参数 alpha 和 lamb）
grads = gradfilter_ema(model, grads=grads, alpha=alpha, lamb=lamb)
### 选项 2：Grokfast-MA（有参数 window_size 和 lamb）
# grads = gradfilter_ma(model, grads=grads, window_size=window_size, lamb=lamb)

optimizer.step() # 调用优化器。
# ... 日志记录及其他代码。

完成！

### 导入
from collections import deque
from typing import Dict, Optional, Literal
import torch
import torch.nn as nn


### Grokfast
def gradfilter_ema(
    m: nn.Module,
    grads: Optional[Dict[str, torch.Tensor]] = None,
    alpha: float = 0.99,
    lamb: float = 5.0,
) -> Dict[str, torch.Tensor]:
    if grads is None:
        grads = {n: p.grad.data.detach() for n, p in m.named_parameters() if p.requires_grad}

    for n, p in m.named_parameters():
        if p.requires_grad:
            grads[n] = grads[n] * alpha + p.grad.data.detach() * (1 - alpha)
            p.grad.data = p.grad.data + grads[n] * lamb

    return grads


### Grokfast-MA
def gradfilter_ma(
    m: nn.Module,
    grads: Optional[Dict[str, deque]] = None,
    window_size: int = 128,
    lamb: float = 5.0,
    filter_type: Literal['mean', 'sum'] = 'mean',
    warmup: bool = True,
    trigger: bool = False,
) -> Dict[str, deque]:
    if grads is None:
        grads = {n: deque(maxlen=window_size) for n, p in m named parameters() if p requires grad}

    for n, p in m named parameters():
        if p requires grad:
            grads[n].append(p.grad.data.detach())

            if not warmup or len(grads[n]) == window_size and not trigger:
                if filter_type == "mean":
                    avg = sum(grads[n]) / len(grads[n])
                elif filter_type == "sum":
                    avg = sum(grads[n])
                else:
                    raise ValueError(f"Unrecognized filter_type {filter_type}")
                p.grad.data = p.grad.data + avg * lamb

    return grads

参数说明

1. Grokfast (gradfilter_ema)

   • m: nn.Module: 包含所有可训练参数的模型。
   • grads: Optional[Dict[str, torch.Tensor]] = None: 运行时内存（EMA）。初始值设为 None。后续递归地传入该方法的输出。
   • alpha: float = 0.98: EMA 的动量超参数。
   • lamb: float = 2.0: 滤波器的放大因子超参数。

2. Grokfast-MA (gradfilter_ma)

   • m: nn.Module: 包含所有可训练参数的模型。
   • grads: Optional[Dict[str, deque]] = None: 运行时内存（窗口移动平均队列）。初始值设为 None。后续递归地传入该方法的输出。
   • window_size: int = 100: 滤波器窗口的宽度。额外的内存需求会随着窗口大小线性增加。
   • lamb: float = 5.0: 滤波器的放大因子超参数。
   • filter_type: Literal['mean', 'sum'] = 'mean': 对运行队列的聚合方式。
   • warmup: bool = True: 如果为真，则在队列填满之前不应用滤波器。
   • trigger: bool = False: 仅用于消融实验。如果为真，则完全不应用滤波器。

复现

我们还注意到每次运行所需的额外计算资源。时间和内存开销均使用单块 GTX 1080 Ti GPU 进行测量。

安装

这将安装用于预处理数据和汇总结果的附加包。

conda create -n grok python=3.10 && conda activate grok
git clone https://github.com/ironjr/grokfast
pip install -r requirements.txt

算法性数据（Transformer 解码器，Grokfast-MA）

# python main.py --label test # 基线。
python main.py --label test --filter ma --window_size 100 --lamb 5.0 --weight_decay 0.01

算法性数据（Transformer 解码器，Grokfast）

# python main.py --label test # 基线。
python main.py --label test --filter ema --alpha 0.98 --lamb 2.0 --weight_decay 0.005

MNIST（MLP）

# python main_mnist.py --label test # 基线。
python main_mnist.py --label test --alpha 0.8 --lamb 0.1 --weight_decay 2.0

IMDb（LSTM）

• 在训练之前，请从 Google Drive 或 百度网盘（提取码：vdp7）下载 IMDb 数据集。

# python main_imdb.py --label test # 基线。
python main_imdb.py --label test --alpha 0.98 --lamb 2.0 --weight_decay 10.0

QM9（G-CNN）

# python main_qm9.py --label test # 基线。
python main_qm9.py --label test --alpha 0.9 --lamb 1.0 --weight_decay 0.01

常见问题解答

如何选择合适的超参数

这些建议基于我在主论文中展示的实验中的经验。它们可能并不适用于所有问题，也可能有更智能的方法能够取得更好的效果。因此，请将这些建议视为设计您自己的滤波器时的一种可能的起点指南。

1. 截止参数：本研究使用 MA/EMA 滤波器来实现滤波技术。对于 MA 滤波器，截止频率由窗口大小决定；对于 EMA 滤波器，截止频率由动量参数决定。
   1. 大致确定您希望达到的加速倍数。 例如，在主论文中，截止参数是根据原始的 grokking 报告确定的，该报告表明泛化发生的速度比过拟合慢约 100 倍。因此，我们希望实现 N=100 倍的加速。
   2. 设定截止参数搜索的基准值。 对于 MA，我从窗口大小 w=N=100 开始；对于 EMA，我从满足 alpha^{N} = alpha^{100} = 0.1 的动量参数 alpha 开始（大约为 0.98）。
   3. 在基准值附近进行超参数搜索。 我围绕 (1.b) 中设定的值进行了超参数扫描。
2. 权重衰减：权重衰减通常在优化器构造函数中设置（例如，optimizer = optim.Adam(m.parameters(), weight_decay=wd)）。
   1. 从该任务的默认权重衰减开始。 例如，该任务中最常用的 GitHub 仓库所选择的值。
   2. 固定权重衰减，首先尝试找到 Grokfast 滤波器参数（动量、窗口大小和幅度）的最佳设置。 尽管权重衰减确实会影响最佳滤波器参数的取值，但根据我的经验，其影响似乎并不显著。
   3. 逐步增加权重衰减的值。 从 X1 开始，然后尝试 (X2, X5, X10)。我未能通过将权重衰减提高到默认值的 100 倍而获得更好的结果。

致谢

我们的代码大量借鉴了以下项目：

• Ziming Liu 等人，“Omnigrok：超越算法性数据的 grokking”，ICLR 2023。[arXiv] [代码]
• Alethea Power 等人，“Grokking：在小型算法性数据集上超越过拟合的泛化”，arXiv 预印本 arXiv:2201.02177。[arXiv] [代码]
• @danielmamay 对 Grokking 的重新实现。[代码]

感谢大家提供的有用参考！

引用

如果您觉得我们的项目有用，请引用我们！

@article{lee2024grokfast,
    title={{Grokfast}: 加速 grokking 通过放大缓慢梯度},
    author={Lee, Jaerin and Kang, Bong Gyun and Kim, Kihoon and Lee, Kyoung Mu},
    journal={arXiv preprint arXiv:2405.20233},
    year={2024}
}

星标历史

联系方式

如有任何问题，请发送邮件至 jarin.lee@gmail.com。

Grokfast 快速上手指南

Grokfast 是一个旨在加速机器学习中“顿悟”（Grokking）现象的优化器增强工具。它通过放大参数梯度中的低频分量（即缓慢变化但能诱导泛化的成分），显著缩短模型从过拟合到泛化的等待时间，适用于图像、语言和图结构等多种任务。

环境准备

• 操作系统: Linux, macOS, Windows
• Python 版本: 推荐 Python 3.8+ (实验复现环境为 Python 3.10)
• 核心依赖:
  • PyTorch (必须)
  • 无需其他额外第三方库即可运行核心功能。

国内加速建议： 安装 PyTorch 时，推荐使用清华或中科大镜像源以加快下载速度：

pip install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

Grokfast 的核心功能仅需一个 Python 文件，无需通过 pip 安装包。

1. 下载核心脚本 直接从 GitHub 仓库下载 grokfast.py 文件到你的项目目录：

   wget https://raw.githubusercontent.com/ironjr/grokfast/main/grokfast.py
   

   (如果无法访问 GitHub，可手动下载该文件并保存为 grokfast.py)

2. (可选) 安装复现实验依赖 如果你需要运行官方提供的复现代码（包含数据预处理和结果汇总），请执行：

   git clone https://github.com/ironjr/grokfast
   cd grokfast
   pip install -r requirements.txt
   

基本使用

只需在现有的训练循环中插入几行代码，即可启用 Grokfast。

1. 导入辅助函数

在你的训练脚本中导入以下函数：

from grokfast import gradfilter_ma, gradfilter_ema

2. 初始化状态变量

在训练循环开始前，初始化梯度记忆变量：

grads = None

3. 修改训练循环

在 loss.backward() 之后、optimizer.step() 之前，插入过滤函数。

方案 A：使用 Grokfast-EMA (指数移动平均，推荐用于大多数场景)

# ... 在优化循环中
loss.backward()  # 计算梯度

# 应用 Grokfast-EMA 过滤器
# alpha: 动量超参数 (默认 0.98), lamb: 放大系数 (默认 2.0)
grads = gradfilter_ema(model, grads=grads, alpha=0.98, lamb=2.0)

optimizer.step()  # 更新参数

方案 B：使用 Grokfast-MA (滑动窗口平均)

# ... 在优化循环中
loss.backward()  # 计算梯度

# 应用 Grokfast-MA 过滤器
# window_size: 窗口大小 (默认 100), lamb: 放大系数 (默认 5.0)
grads = gradfilter_ma(model, grads=grads, window_size=100, lamb=5.0)

optimizer.step()  # 更新参数

参数调优建议

• alpha (EMA): 控制截止频率。若希望加速倍数约为 $N$，可尝试设定 $\alpha$ 使得 $\alpha^N \approx 0.1$（例如加速 100 倍，$\alpha \approx 0.98$）。
• window_size (MA): 对应加速倍数的粗略估计值。
• lamb: 梯度放大系数，通常设置在 1.0 到 5.0 之间，需配合具体任务调整。
• weight_decay: 建议保留优化器中原有的权重衰减设置，先固定它再调整 Grokfast 参数。