Group QKAN for Transformers
To demonstrate the generative capabilities of the QKAN model, we adapt the excellent karpathy/nanoGPT implementation and replace the standard MLP components with QKAN-based modules. This variant is referred to as Kansformer.
Note: Although our main paper introduces a different method, the Kansformer variant shown here provides a much more efficient implementation.
Motivation
Previously, we introduced Hybrid QKAN (HQKAN) — a variant that retains the expressive power of QKAN while using significantly fewer parameters than standard MLPs.
While HQKAN addresses the parameter scalability challenge of KAN-family models, a new issue arises as the model grows in size:
Having too many learnable variational activation functions (VAFs) in a single layer can make optimization difficult and unstable.
To solve this, X. Yang et al., in “Kolmogorov–Arnold Transformer”, proposed grouping the VAFs within each layer. In this approach, VAFs are shared across groups, which reduces the number of independent parameters and improves training efficiency.
Our Implementation
We implement the GPT-2 architecture with Hybrid Grouped QKAN (HG-QKAN) modules and train the model on the WebText dataset.
Note: However, if we have already used Hybrid architecture to reduce the number of VAFs, grouping may not be necessary. In our previous experiments, we found that HQKAN without grouping performed better than HG-QKAN, if the number of VAFs was not significantly large.
[1]:
import inspect
import math
import os
import pickle
import time
import warnings
from ast import literal_eval
from collections import defaultdict
from dataclasses import dataclass
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.utils.data import Dataset
from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm
from transformers import GPT2Tokenizer
from qkan import QKAN
warnings.filterwarnings("ignore", category=FutureWarning)
device = "cuda"
[2]:
# Define dataset
class WebTextDataset(Dataset):
"""
WebText Dataset
"""
def __init__(self, split, model_type, block_size=1024, vocab_size=50257):
assert split in {"train", "test", "valid"}
self.split = split
self.block_size = block_size
self.vocab_size = vocab_size
self.tokenizer = GPT2Tokenizer.from_pretrained(model_type)
self.tokenized_dataset_path = f"./datasets/webtext/webtext.{split}.pkl"
if not os.path.isfile(self.tokenized_dataset_path):
self.tokenized_dataset = []
self.json_path = f"./datasets/webtext/webtext.{split}.jsonl"
assert os.path.isfile(self.json_path)
self.data = pd.read_json(path_or_buf=self.json_path, lines=True)
tokenized_data = []
tokenized_lengths = []
for _, row in tqdm(
self.data.iterrows(), desc="Tokenizing", total=len(self.data)
):
text = row["text"]
tokenized = self.tokenizer.encode(
text=text, add_special_tokens=False
)
tokenized_length = len(tokenized)
tokenized_data.append(tokenized)
tokenized_lengths.append(tokenized_length)
self.tokenized_dataset += tokenized
with open(self.tokenized_dataset_path, "wb") as f:
pickle.dump(self.tokenized_dataset, f)
with open(self.tokenized_dataset_path, "rb") as f:
self.tokenized_dataset = pickle.load(f)
def __len__(self):
return len(self.tokenized_dataset) - 2 * self.block_size
def get_vocab_size(self):
return self.vocab_size
def get_block_size(self):
return self.block_size
def __getitem__(self, idx):
x = self.tokenized_dataset[idx : idx + self.block_size]
y = self.tokenized_dataset[
idx + self.block_size : idx + 2 * self.block_size
]
assert len(x) == self.block_size, f"Unexpected len: {len(x)}"
assert len(y) == self.block_size, f"Unexpected len: {len(y)}"
x = torch.tensor(x, dtype=torch.long)
y = torch.tensor(y, dtype=torch.long)
return x, y
[ ]:
class LayerNorm(nn.Module):
"""LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False"""
def __init__(self, ndim, bias):
super().__init__()
self.weight = nn.Parameter(torch.ones(ndim))
self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
def forward(self, input):
return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
class CN:
"""a lightweight configuration class inspired by yacs"""
def __init__(self, **kwargs):
self.__dict__.update(kwargs)
def __str__(self):
return self._str_helper(0)
def _str_helper(self, indent):
"""need to have a helper to support nested indentation for pretty printing"""
parts = []
for k, v in self.__dict__.items():
if isinstance(v, CN):
parts.append("%s:\n" % k)
parts.append(v._str_helper(indent + 1))
else:
parts.append("%s: %s\n" % (k, v))
parts = [" " * (indent * 4) + p for p in parts]
return "".join(parts)
def to_dict(self):
"""return a dict representation of the config"""
return {
k: v.to_dict() if isinstance(v, CN) else v for k, v in self.__dict__.items()
}
def merge_from_dict(self, d):
self.__dict__.update(d)
def merge_from_args(self, args):
"""
update the configuration from a list of strings that is expected
to come from the command line, i.e. sys.argv[1:].
The arguments are expected to be in the form of `--arg=value`, and
the arg can use . to denote nested sub-attributes. Example:
--model.n_layer=10 --trainer.batch_size=32
"""
for arg in args:
keyval = arg.split("=")
assert len(keyval) == 2, (
"expecting each override arg to be of form --arg=value, got %s" % arg
)
key, val = keyval # unpack
# first translate val into a python object
try:
val = literal_eval(val)
"""
need some explanation here.
- if val is simply a string, literal_eval will throw a ValueError
- if val represents a thing (like an 3, 3.14, [1,2,3], False, None, etc.) it will get created
"""
except ValueError:
pass
# find the appropriate object to insert the attribute into
assert key[:2] == "--"
key = key[2:] # strip the '--'
keys = key.split(".")
obj = self
for k in keys[:-1]:
obj = getattr(obj, k)
leaf_key = keys[-1]
# ensure that this attribute exists
assert hasattr(
obj, leaf_key
), f"{key} is not an attribute that exists in the config"
# overwrite the attribute
print("command line overwriting config attribute %s with %s" % (key, val))
setattr(obj, leaf_key, val)
class CausalSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
assert config.n_embd % config.n_head == 0
# key, query, value projections for all heads, but in a batch
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
# output projection
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
# regularization
self.attn_dropout = nn.Dropout(config.dropout)
self.resid_dropout = nn.Dropout(config.dropout)
self.n_head = config.n_head
self.n_embd = config.n_embd
self.dropout = config.dropout
# flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
self.flash = hasattr(torch.nn.functional, "scaled_dot_product_attention")
if not self.flash:
print(
"WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0"
)
# causal mask to ensure that attention is only applied to the left in the input sequence
self.register_buffer(
"bias",
torch.tril(torch.ones(config.block_size, config.block_size)).view(
1, 1, config.block_size, config.block_size
),
)
def forward(self, x):
B, T, C = (
x.size()
) # batch size, sequence length, embedding dimensionality (n_embd)
# calculate query, key, values for all heads in batch and move head forward to be the batch dim
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
k = k.view(B, T, self.n_head, C // self.n_head).transpose(
1, 2
) # (B, nh, T, hs)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(
1, 2
) # (B, nh, T, hs)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(
1, 2
) # (B, nh, T, hs)
# causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
if self.flash:
# efficient attention using Flash Attention CUDA kernels
y = torch.nn.functional.scaled_dot_product_attention(
q,
k,
v,
attn_mask=None,
dropout_p=self.dropout if self.training else 0,
is_causal=True,
)
else:
# manual implementation of attention
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float("-inf"))
att = F.softmax(att, dim=-1)
att = self.attn_dropout(att)
y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
y = (
y.transpose(1, 2).contiguous().view(B, T, C)
) # re-assemble all head outputs side by side
# output projection
y = self.resid_dropout(self.c_proj(y))
return y
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
self.attn = CausalSelfAttention(config)
self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
els = math.ceil(math.log2(config.n_embd))
self.mlp = nn.ModuleDict(
dict(
c=nn.Sequential(
nn.Linear(config.n_embd, els),
QKAN(
width=[els, els],
reps=1,
group=3,
preact_trainable=True,
postact_weight_trainable=True,
postact_bias_trainable=True,
ba_trainable=True,
device=device,
),
nn.Linear(els, config.n_embd),
),
dropout=nn.Dropout(config.dropout),
)
)
m = self.mlp
self.mlpf = lambda x: m.dropout(m.c(x))
def forward(self, x):
x = x + self.attn(self.ln_1(x))
x = x + self.mlpf(self.ln_2(x))
return x
@dataclass
class GPTConfig:
block_size: int = 1024
vocab_size: int = (
50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
)
n_layer: int = 12
n_head: int = 12
n_embd: int = 768
dropout: float = 0.0
bias: bool = (
True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
)
class GPT(nn.Module):
def __init__(self, config):
super().__init__()
assert config.vocab_size is not None
assert config.block_size is not None
self.config = config
self.transformer = nn.ModuleDict(
dict(
wte=nn.Embedding(config.vocab_size, config.n_embd),
wpe=nn.Embedding(config.block_size, config.n_embd),
drop=nn.Dropout(config.dropout),
h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
ln_f=LayerNorm(config.n_embd, bias=config.bias),
)
)
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
# with weight tying when using torch.compile() some warnings get generated:
# "UserWarning: functional_call was passed multiple values for tied weights.
# This behavior is deprecated and will be an error in future versions"
# not 100% sure what this is, so far seems to be harmless. TODO investigate
self.transformer.wte.weight = (
self.lm_head.weight
) # https://paperswithcode.com/method/weight-tying
# init all weights
self.apply(self._init_weights)
# apply special scaled init to the residual projections, per GPT-2 paper
for pn, p in self.named_parameters():
if pn.endswith("c_proj.weight"):
torch.nn.init.normal_(
p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer)
)
# report number of parameters
print("number of parameters: %.2fM" % (self.get_num_params() / 1e6,))
def get_num_params(self, non_embedding=True):
"""
Return the number of parameters in the model.
For non-embedding count (default), the position embeddings get subtracted.
The token embeddings would too, except due to the parameter sharing these
params are actually used as weights in the final layer, so we include them.
"""
n_params = sum(p.numel() for p in self.parameters())
if non_embedding:
n_params -= self.transformer.wpe.weight.numel()
return n_params
def _init_weights(self, module):
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
def forward(self, idx, targets=None):
device = idx.device
b, t = idx.size()
assert (
t <= self.config.block_size
), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
# forward the GPT model itself
tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
x = self.transformer.drop(tok_emb + pos_emb)
for block in self.transformer.h:
x = block(x)
x = self.transformer.ln_f(x)
if targets is not None:
# if we are given some desired targets also calculate the loss
logits = self.lm_head(x)
loss = F.cross_entropy(
logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1
)
else:
# inference-time mini-optimization: only forward the lm_head on the very last position
logits = self.lm_head(
x[:, [-1], :]
) # note: using list [-1] to preserve the time dim
loss = None
return logits, loss
def crop_block_size(self, block_size):
# model surgery to decrease the block size if necessary
# e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
# but want to use a smaller block size for some smaller, simpler model
assert block_size <= self.config.block_size
self.config.block_size = block_size
self.transformer.wpe.weight = nn.Parameter(
self.transformer.wpe.weight[:block_size]
)
for block in self.transformer.h:
if hasattr(block.attn, "bias"):
block.attn.bias = block.attn.bias[:, :, :block_size, :block_size]
@classmethod
def from_pretrained(cls, model_type, override_args=None):
assert model_type in {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl"}
override_args = override_args or {} # default to empty dict
# only dropout can be overridden see more notes below
assert all(k == "dropout" for k in override_args)
from transformers import GPT2LMHeadModel
print("loading weights from pretrained gpt: %s" % model_type)
# n_layer, n_head and n_embd are determined from model_type
config_args = {
"gpt2": dict(n_layer=12, n_head=12, n_embd=768), # 124M params
"gpt2-medium": dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
"gpt2-large": dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
"gpt2-xl": dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
}[model_type]
print("forcing vocab_size=50257, block_size=1024, bias=True")
config_args["vocab_size"] = 50257 # always 50257 for GPT model checkpoints
config_args["block_size"] = 1024 # always 1024 for GPT model checkpoints
config_args["bias"] = True # always True for GPT model checkpoints
# we can override the dropout rate, if desired
if "dropout" in override_args:
print(f"overriding dropout rate to {override_args['dropout']}")
config_args["dropout"] = override_args["dropout"]
# create a from-scratch initialized minGPT model
config = GPTConfig(**config_args)
model = GPT(config)
sd = model.state_dict()
sd_keys = sd.keys()
sd_keys = [
k for k in sd_keys if not k.endswith(".attn.bias")
] # discard this mask / buffer, not a param
# init a huggingface/transformers model
model_hf = GPT2LMHeadModel.from_pretrained(model_type)
sd_hf = model_hf.state_dict()
# copy while ensuring all of the parameters are aligned and match in names and shapes
sd_keys_hf = sd_hf.keys()
sd_keys_hf = [
k for k in sd_keys_hf if not k.endswith(".attn.masked_bias")
] # ignore these, just a buffer
sd_keys_hf = [
k for k in sd_keys_hf if not k.endswith(".attn.bias")
] # same, just the mask (buffer)
transposed = [
"attn.c_attn.weight",
"attn.c_proj.weight",
"mlp.c_fc.weight",
"mlp.c_proj.weight",
]
# basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
# this means that we have to transpose these weights when we import them
assert len(sd_keys_hf) == len(
sd_keys
), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
for k in sd_keys_hf:
if any(k.endswith(w) for w in transposed):
# special treatment for the Conv1D weights we need to transpose
assert sd_hf[k].shape[::-1] == sd[k].shape
with torch.no_grad():
sd[k].copy_(sd_hf[k].t())
else:
# vanilla copy over the other parameters
assert sd_hf[k].shape == sd[k].shape
with torch.no_grad():
sd[k].copy_(sd_hf[k])
return model
def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
# start with all of the candidate parameters
param_dict = {pn: p for pn, p in self.named_parameters()}
# filter out those that do not require grad
param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
# create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
# i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
optim_groups = [
{"params": decay_params, "weight_decay": weight_decay},
{"params": nodecay_params, "weight_decay": 0.0},
]
num_decay_params = sum(p.numel() for p in decay_params)
num_nodecay_params = sum(p.numel() for p in nodecay_params)
print(
f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters"
)
print(
f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters"
)
# Create AdamW optimizer and use the fused version if it is available
fused_available = "fused" in inspect.signature(torch.optim.AdamW).parameters
use_fused = fused_available and device_type == "cuda"
extra_args = dict(fused=True) if use_fused else dict()
optimizer = torch.optim.AdamW(
optim_groups, lr=learning_rate, betas=betas, **extra_args
)
print(f"using fused AdamW: {use_fused}")
return optimizer
def estimate_mfu(self, fwdbwd_per_iter, dt):
"""estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS"""
# first estimate the number of flops we do per iteration.
# see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
N = self.get_num_params()
cfg = self.config
L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd // cfg.n_head, cfg.block_size
flops_per_token = 6 * N + 12 * L * H * Q * T
flops_per_fwdbwd = flops_per_token * T
flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
# express our flops throughput as ratio of A100 bfloat16 peak flops
flops_achieved = flops_per_iter * (1.0 / dt) # per second
flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
mfu = flops_achieved / flops_promised
return mfu
@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
"""
Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
the sequence max_new_tokens times, feeding the predictions back into the model each time.
Most likely you'll want to make sure to be in model.eval() mode of operation for this.
"""
for _ in range(max_new_tokens):
# if the sequence context is growing too long we must crop it at block_size
idx_cond = (
idx
if idx.size(1) <= self.config.block_size
else idx[:, -self.config.block_size :]
)
# forward the model to get the logits for the index in the sequence
logits, _ = self(idx_cond)
# pluck the logits at the final step and scale by desired temperature
logits = logits[:, -1, :] / temperature
# optionally crop the logits to only the top k options
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float("Inf")
# apply softmax to convert logits to (normalized) probabilities
probs = F.softmax(logits, dim=-1)
# sample from the distribution
idx_next = torch.multinomial(probs, num_samples=1)
# append sampled index to the running sequence and continue
idx = torch.cat((idx, idx_next), dim=1)
return idx
[4]:
def metrics(y, y_pred):
"""
y: (B, T) INT - True labels
y_pred: (B, T, C) FLOAT - Predicted probabilities
Returns:
- Perplexity
- F1 Score
- Precision
- Recall
- Cross Entropy
"""
# Make sure y_pred is between 0 and 1
if not (np.all(y_pred >= 0) and np.all(y_pred <= 1)):
# Softmax
y_pred = np.exp(y_pred) / np.sum(np.exp(y_pred), axis=-1, keepdims=True)
assert np.all(y_pred >= 0) and np.all(y_pred <= 1), "y_pred must be between 0 and 1"
# Add a small epsilon for numerical stability
epsilon = 1e-9
y_pred = np.clip(y_pred, epsilon, 1 - epsilon)
# Cross Entropy
y_one_hot = np.eye(y_pred.shape[-1])[y]
cross_entropy = -np.mean(np.sum(y_one_hot * np.log(y_pred), axis=-1))
# Perplexity
perplexity = 2**cross_entropy
# Predicted classes
y_pred_class = np.argmax(y_pred, axis=-1)
# True Positives, False Positives, and False Negatives
TP = np.sum(y == y_pred_class)
FP = np.sum(y != y_pred_class)
FN = FP # Binary setup, false positives and false negatives are equivalent
# Precision, Recall
precision = TP / (TP + FP)
recall = TP / (TP + FN)
# F1 Score
f1 = 2 * (precision * recall) / (precision + recall)
return perplexity, f1, precision, recall, cross_entropy
def eval_split(
trainer, split, max_batches, batch_size, model, train_dataset, test_dataset
):
dataset = {"train": train_dataset, "test": test_dataset}[split]
results = []
loader = DataLoader(dataset, batch_size=batch_size, num_workers=0, drop_last=False)
for b, (x, y) in enumerate(loader):
x = x.to(trainer.device)
y = y.to(trainer.device)
block_size = y.shape[1]
logits, loss = model(x, y)
probs = F.softmax(logits, dim=-1)
_, y_pred = torch.topk(probs, k=block_size, dim=-1)
perplexity, f1, precision, recall, cross_entropy = metrics(
y=y.cpu().numpy(), y_pred=probs.cpu().numpy()
)
results.append((loss, perplexity, f1, precision, recall, cross_entropy))
if max_batches is not None and b + 1 >= max_batches:
break
rt = torch.tensor(results, dtype=torch.float)
print("%s loss: %.2f" % (split, rt.mean(dim=0)[0]))
return rt.mean(dim=0)
[5]:
Dataset = WebTextDataset
test_dataset = Dataset("test", "gpt2")
train_dataset = Dataset("train", "gpt2")
print("test_dataset: ", len(test_dataset))
print("train_dataset: ", len(train_dataset))
# create a GPT instance
model_config = GPTConfig()
model_config.model_type = "gpt2"
model_config.vocab_size = train_dataset.get_vocab_size()
model_config.block_size = train_dataset.get_block_size()
model = GPT(model_config)
print(
"Number of parameters: %.2fM"
% (sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6,)
)
test_dataset: 2898230
train_dataset: 146323009
number of parameters: 67.19M
Number of parameters: 67.97M
[6]:
class Trainer:
@staticmethod
def get_default_config():
C = CN()
# device to train on
C.device = "auto"
# dataloder parameters
C.num_workers = 4
# optimizer parameters
C.max_iters = None
C.batch_size = 64
C.learning_rate = 3e-4
C.betas = (0.9, 0.95)
C.weight_decay = 0.1 # only applied on matmul weights
C.grad_norm_clip = 1.0
return C
def __init__(self, config, model, train_dataset):
self.config = config
self.model = model
self.optimizer = None
self.train_dataset = train_dataset
self.callbacks = defaultdict(list)
# determine the device we'll train on
if config.device == "auto":
self.device = "cuda" if torch.cuda.is_available() else "cpu"
else:
self.device = config.device
self.model = self.model.to(self.device)
print("running on device", self.device)
# variables that will be assigned to trainer class later for logging and etc
self.iter_num = 0
self.iter_time = 0.0
self.iter_dt = 0.0
def add_callback(self, onevent: str, callback):
self.callbacks[onevent].append(callback)
def set_callback(self, onevent: str, callback):
self.callbacks[onevent] = [callback]
def trigger_callbacks(self, onevent: str):
for callback in self.callbacks.get(onevent, []):
callback(self)
def run(self):
model, config = self.model, self.config
# setup the optimizer
self.optimizer = model.configure_optimizers(config.weight_decay, config.learning_rate, config.betas, self.device)
# setup the dataloader
train_loader = DataLoader(
self.train_dataset,
sampler=torch.utils.data.RandomSampler(
self.train_dataset, replacement=True, num_samples=int(1e10)
),
shuffle=False,
pin_memory=True,
batch_size=config.batch_size,
num_workers=config.num_workers,
)
model.train()
self.iter_num = 0
self.iter_time = time.time()
data_iter = iter(train_loader)
while True:
# fetch the next batch (x, y) and re-init iterator if needed
try:
batch = next(data_iter)
except StopIteration:
data_iter = iter(train_loader)
batch = next(data_iter)
batch = [t.to(self.device) for t in batch]
x, y = batch
# forward the model
logits, self.loss = model(x, y)
# backprop and update the parameters
model.zero_grad(set_to_none=True)
self.loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_norm_clip)
self.optimizer.step()
self.trigger_callbacks("on_batch_end")
self.iter_num += 1
tnow = time.time()
self.iter_dt = tnow - self.iter_time
self.iter_time = tnow
# termination conditions
if config.max_iters is not None and self.iter_num >= config.max_iters:
break
[7]:
# create a Trainer object
train_config = Trainer.get_default_config()
train_config.learning_rate = 3e-4
train_config.max_iters = 2000
train_config.num_workers = 0
train_config.batch_size = 1
train_config.device = device
trainer = Trainer(train_config, model, train_dataset)
running on device cuda
[8]:
def batch_end_callback(trainer):
if trainer.iter_num % 100 == 0:
print(f"iter_dt {trainer.iter_dt * 1000:.2f}ms; iter {trainer.iter_num}: train loss {trainer.loss.item():.5f}")
trainer.set_callback('on_batch_end', batch_end_callback)
trainer.run()
num decayed parameter tensors: 122, with 67,885,824 parameters
num non-decayed parameter tensors: 98, with 84,600 parameters
using fused AdamW: True
iter_dt 0.00ms; iter 0: train loss 10.95520
iter_dt 37.84ms; iter 100: train loss 8.03499
iter_dt 37.37ms; iter 200: train loss 8.12177
iter_dt 37.87ms; iter 300: train loss 7.91832
iter_dt 37.83ms; iter 400: train loss 8.11329
iter_dt 37.46ms; iter 500: train loss 8.05757
iter_dt 37.92ms; iter 600: train loss 7.64972
iter_dt 37.87ms; iter 700: train loss 7.60039
iter_dt 37.85ms; iter 800: train loss 7.94277
iter_dt 37.84ms; iter 900: train loss 7.85295
iter_dt 37.84ms; iter 1000: train loss 10.16237
iter_dt 37.88ms; iter 1100: train loss 8.37014
iter_dt 37.87ms; iter 1200: train loss 7.67977
iter_dt 37.90ms; iter 1300: train loss 8.22389
iter_dt 37.85ms; iter 1400: train loss 8.00896
iter_dt 37.92ms; iter 1500: train loss 9.33540
iter_dt 37.94ms; iter 1600: train loss 7.94113
iter_dt 37.47ms; iter 1700: train loss 8.19066
iter_dt 37.93ms; iter 1800: train loss 7.85500
iter_dt 37.92ms; iter 1900: train loss 7.81131
[9]:
model.eval()
with torch.no_grad():
train_score = eval_split(
trainer,
"train",
max_batches=5,
batch_size=1,
model=model,
train_dataset=train_dataset,
test_dataset=test_dataset,
)
test_score = eval_split(
trainer,
"test",
max_batches=5,
batch_size=1,
model=model,
train_dataset=train_dataset,
test_dataset=test_dataset,
)
(
train_loss,
train_perplexity,
train_f1,
train_precision,
train_recall,
train_cross_entropy,
) = train_score
(
test_loss,
test_perplexity,
test_f1,
test_precision,
test_recall,
test_cross_entropy,
) = test_score
print("train_loss: ", train_loss)
print("train_perplexity: ", train_perplexity)
print("train_f1: ", train_f1)
print("train_precision: ", train_precision)
print("train_recall: ", train_recall)
print("train_cross_entropy: ", train_cross_entropy)
print("test_loss: ", test_loss)
print("test_perplexity: ", test_perplexity)
print("test_f1: ", test_f1)
print("test_precision: ", test_precision)
print("test_recall: ", test_recall)
print("test_cross_entropy: ", test_cross_entropy)
train loss: 7.63
test loss: 7.59
train_loss: tensor(7.6327)
train_perplexity: tensor(198.4565)
train_f1: tensor(0.0475)
train_precision: tensor(0.0475)
train_recall: tensor(0.0475)
train_cross_entropy: tensor(7.6327)
test_loss: tensor(7.5937)
test_perplexity: tensor(193.1681)
test_f1: tensor(0.0312)
test_precision: tensor(0.0312)
test_recall: tensor(0.0312)
test_cross_entropy: tensor(7.5937)
Compare to MLP
[10]:
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
self.gelu = nn.GELU()
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
self.dropout = nn.Dropout(config.dropout)
def forward(self, x):
x = self.c_fc(x)
x = self.gelu(x)
x = self.c_proj(x)
x = self.dropout(x)
return x
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
self.attn = CausalSelfAttention(config)
self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
self.mlp = MLP(config)
def forward(self, x):
x = x + self.attn(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return x
[11]:
# create a GPT instance
model_config = GPTConfig()
model_config.model_type = "gpt2"
model_config.vocab_size = train_dataset.get_vocab_size()
model_config.block_size = train_dataset.get_block_size()
model = GPT(model_config)
print(
"Number of parameters: %.2fM"
% (sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6,)
)
number of parameters: 123.65M
Number of parameters: 124.44M
[12]:
trainer = Trainer(train_config, model, train_dataset)
trainer.set_callback('on_batch_end', batch_end_callback)
trainer.run()
running on device cuda
num decayed parameter tensors: 50, with 124,318,464 parameters
num non-decayed parameter tensors: 98, with 121,344 parameters
using fused AdamW: True
iter_dt 0.00ms; iter 0: train loss 10.96399
iter_dt 38.10ms; iter 100: train loss 7.46826
iter_dt 38.00ms; iter 200: train loss 7.66300
iter_dt 38.00ms; iter 300: train loss 7.95232
iter_dt 37.63ms; iter 400: train loss 7.66633
iter_dt 37.67ms; iter 500: train loss 8.46683
iter_dt 38.03ms; iter 600: train loss 7.43420
iter_dt 37.69ms; iter 700: train loss 7.37517
iter_dt 37.76ms; iter 800: train loss 8.08870
iter_dt 37.90ms; iter 900: train loss 7.54975
iter_dt 37.37ms; iter 1000: train loss 7.96053
iter_dt 38.32ms; iter 1100: train loss 7.61653
iter_dt 37.92ms; iter 1200: train loss 7.84077
iter_dt 37.52ms; iter 1300: train loss 7.86578
iter_dt 37.99ms; iter 1400: train loss 8.13120
iter_dt 37.65ms; iter 1500: train loss 7.99749
iter_dt 38.02ms; iter 1600: train loss 7.91518
iter_dt 37.67ms; iter 1700: train loss 7.84328
iter_dt 37.97ms; iter 1800: train loss 7.89387
iter_dt 37.39ms; iter 1900: train loss 7.96062
[13]:
model.eval()
with torch.no_grad():
train_score = eval_split(
trainer,
"train",
max_batches=5,
batch_size=1,
model=model,
train_dataset=train_dataset,
test_dataset=test_dataset,
)
test_score = eval_split(
trainer,
"test",
max_batches=5,
batch_size=1,
model=model,
train_dataset=train_dataset,
test_dataset=test_dataset,
)
(
train_loss,
train_perplexity,
train_f1,
train_precision,
train_recall,
train_cross_entropy,
) = train_score
(
test_loss,
test_perplexity,
test_f1,
test_precision,
test_recall,
test_cross_entropy,
) = test_score
print("train_loss: ", train_loss)
print("train_perplexity: ", train_perplexity)
print("train_f1: ", train_f1)
print("train_precision: ", train_precision)
print("train_recall: ", train_recall)
print("train_cross_entropy: ", train_cross_entropy)
print("test_loss: ", test_loss)
print("test_perplexity: ", test_perplexity)
print("test_f1: ", test_f1)
print("test_precision: ", test_precision)
print("test_recall: ", test_recall)
print("test_cross_entropy: ", test_cross_entropy)
train loss: 7.70
test loss: 7.67
train_loss: tensor(7.7001)
train_perplexity: tensor(207.9543)
train_f1: tensor(0.0475)
train_precision: tensor(0.0475)
train_recall: tensor(0.0475)
train_cross_entropy: tensor(7.7001)
test_loss: tensor(7.6732)
test_perplexity: tensor(204.1052)
test_f1: tensor(0.0312)
test_precision: tensor(0.0312)
test_recall: tensor(0.0312)
test_cross_entropy: tensor(7.6732)
Results
The HG-QKAN model achieves better performance than the MLP-based baseline, while using:
Only half the number of parameters
The same training time
This result is both surprising and promising, demonstrating the power and efficiency of QKAN-based architectures, especially when enhanced with grouping and hybrid techniques.