KAN-GPT

使用科尔莫戈洛夫-阿诺德网络（KAN）进行语言建模的生成式预训练 Transformer（GPT）的 PyTorch 实现

从 PyPI 安装

pip install kan_gpt

引用

如果您觉得我们的工作有用，请引用我们！

@misc{GANESH2024KANGPT,
  author       = {Aditya Nalgunda Ganesh},
  title        = {KAN-GPT：使用科尔莫戈洛夫-阿诺德网络（KAN）进行语言建模的生成式预训练 Transformer（GPT）的 PyTorch 实现},
  year         = {2024},
  month        = {5月},
  note         = {发布 1.0.0，2024年5月9日},
  url          = {https://github.com/AdityaNG/kan-gpt/}
}

使用方法

请参阅 KAN_GPT.ipynb 和 kan_gpt/prompt.py 以获取使用示例。以下是模型使用的大致流程：

from kan_gpt.model import GPT
from transformers import GPT2Tokenizer

model_config = GPT.get_default_config()
model_config.model_type = "gpt2"
model_config.vocab_size = 50257
model_config.block_size = 1024
model = GPT(model_config)

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

prompt = "班加罗尔通常被描述为"

prompt_encoded = tokenizer.encode(
  text=prompt, add_special_tokens=False
)

x = torch.tensor(prompt_encoded).unsqueeze(0)

model.eval()
y = model.generate(x, 50)  # 采样 50 个 token

result = tokenizer.decode(y[0])

print(result)

# 班加罗尔通常被描述为印度的硅谷。
# 这座城市在过去二十年里经历了快速的发展……

开发环境搭建

# 克隆仓库
git clone https://github.com/AdityaNG/kan-gpt
cd kan-gpt
git pull

# 下载数据集
python3 -m kan_gpt.download_dataset --dataset tinyshakespeare
python3 -m kan_gpt.download_dataset --dataset mnist
python3 -m kan_gpt.download_dataset --dataset webtext

# 安装开发依赖
pip install -r requirements.txt
pip install -e .

训练

使用以下示例脚本来确保一切正常运行：

WANDB_MODE=offline CUDA_VISIBLE_DEVICE="" python3 -m kan_gpt.train --architecture MLP --batch_size 1 --dummy_dataset --device cpu --max_iters 200
WANDB_MODE=offline CUDA_VISIBLE_DEVICE="" python3 -m kan_gpt.train --architecture KAN --batch_size 1 --dummy_dataset --device cpu --max_iters 200

然后可以使用正式的训练脚本：

python -m kan_gpt.train

提示词生成

您可以按照如下方式提示模型生成文本：

python -m kan_gpt.prompt --prompt "班加罗尔通常被描述为" --model_path (checkpoint)

结果

我们在 Tiny Shakespeare 数据集上训练并比较了 KAN-GPT 和等效的 MLP-GPT 模型。我们观察到 KAN-GPT 的表现略优于 MLP-GPT。我们正在计划进一步的实验以深入研究。结果如下所示：

指标

待办事项

• 集成 minGPT 和 pykan
• 为 WebText 编写数据集下载脚本
• 为 WebText 编写 PyTorch 数据集解析器
• 为 tinyshakespeare 编写 PyTorch 数据集解析器
• 为 KAN-GPT 制定小型训练 PoC
  • 集成 KAN.train_kan 中的 KAN 训练逻辑
  • 在不出现内存问题的情况下训练一个示例批次
• 为 MLP-GPT 制定小型训练 PoC
• 将 MLP-GPT 在 webtext 数据集上作为基线进行训练
• 将 KAN-GPT 在 webtext 数据集上作为基线进行训练
• 比较 KAN-GPT 和 MLP-GPT 的各项指标
• 自动保存检查点
• 自动将检查点保存到 W&B
• 自动从 git / huggingface 下载模型权重
• W&B 超参数扫描脚本
• 用于交互模式加载检查点的脚本
• 减少 requirements.txt 中的约束条件
• 定义用于训练和扫描参数的 Pydantic 模型
• 对软件包进行裁剪，移除未使用的代码
• 将训练脚本迁移到 PyTorch Lightning
• 文档：mkdocs gh-deploy
• 与 efficient-kan 集成
• 测试用例
  • KAN：前向-反向测试
  • GPT：前向-反向测试
  • KAN_GPT：前向-反向测试
  • EFFICIENT_KAN：前向-反向测试

开发贡献

请阅读 CONTRIBUTING.md 文件。

参考文献

• minGPT
• pykan
• webtext
• tinyshakespeare

KAN-GPT 快速上手指南

KAN-GPT 是一个基于 PyTorch 的项目，它使用 Kolmogorov-Arnold Networks (KANs) 来实现生成式预训练变压器（GPT）进行语言建模。本指南将帮助你快速安装并运行该模型。

环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python: 建议版本 3.8 及以上
• 核心依赖:
  • PyTorch
  • Transformers (Hugging Face)
  • Weights & Biases (可选，用于训练监控)
• 硬件: 推荐使用 NVIDIA GPU 以加速训练和推理（CPU 亦可运行，但速度较慢）

提示：国内开发者建议使用清华源或阿里源加速 Python 包的安装。

安装步骤

1. 从 PyPI 安装（推荐）

最简单的方式是通过 pip 直接安装发布版：

pip install kan_gpt

国内加速安装：

pip install kan_gpt -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 源码安装（开发模式）

如果你需要修改源码或参与贡献，请克隆仓库并进行可编辑安装：

# 下载代码库
git clone https://github.com/AdityaNG/kan-gpt
cd kan-gpt

# 安装依赖
pip install -r requirements.txt
pip install -e .

国内加速安装依赖：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install -e .

3. 下载数据集（可选）

如果需要本地训练或测试，可以下载内置支持的数据集（如 Tiny Shakespeare, MNIST, WebText）：

python3 -m kan_gpt.download_dataset --dataset tinyshakespeare
python3 -m kan_gpt.download_dataset --dataset mnist
python3 -m kan_gpt.download_dataset --dataset webtext

基本使用

以下是最简单的文本生成示例。该示例加载默认配置的 GPT 模型，并使用 GPT-2 的分词器对提示词进行编码，生成后续文本。

import torch
from kan_gpt.model import GPT
from transformers import GPT2Tokenizer

# 1. 初始化模型配置
model_config = GPT.get_default_config()
model_config.model_type = "gpt2"
model_config.vocab_size = 50257
model_config.block_size = 1024
model = GPT(model_config)

# 2. 加载分词器
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 3. 准备提示词
prompt = "Bangalore is often described as the "

# 4. 编码输入
prompt_encoded = tokenizer.encode(
  text=prompt, add_special_tokens=False
)

x = torch.tensor(prompt_encoded).unsqueeze(0)

# 5. 设置模型为评估模式并生成文本
model.eval()
# generate 第二个参数表示生成的 token 数量
y = model.generate(x, 50)  

# 6. 解码并输出结果
result = tokenizer.decode(y[0])

print(result)

预期输出示例：

Bangalore is often described as the Silicon Valley of India.
The city has witnessed rapid growth in the past two decades.....

命令行快速推理

如果你已经训练好了模型或拥有检查点文件，可以直接通过命令行进行交互：

python -m kan_gpt.prompt --prompt "Bangalore is often described as the " --model_path <你的检查点路径>