Standard practice when training transformers (GPT, BERT, ViT) is to apply weight decay to weight matrices but not to biases, normalisation parameters or embeddings (decaying a LayerNorm scale towards 0 just hurts).
Given a dict of named parameters (NumPy arrays), implement:
split_param_groups(params, no_decay_keywords=("bias", "norm", "embed")): return (decay, no_decay), two sorted lists of names. A parameter is not decayed if it has fewer than 2 dimensions or its name contains any of the keywords (case-insensitive).count_params(params, names): the total number of elements in those parametersapply_decay(params, decay_names, lr, wd): return a new dict where only the decayed parameters are multiplied by (1 - lr * wd) (decoupled, AdamW-style), leaving the inputs unchanged