nGPT-pytorch

nGPT (normalized GPT) - Pytorch

Quick implementation of nGPT, learning entirely on the hypersphere, from NvidiaAI. The question is whether there is any loss of expressivity they swept under the rug, but I'll take it with good faith.

This type of network should also be studied in the context of continual learning and loss of plasticity

Adaptation to vision transformers is here

Update: The normalized feedforward was successfully applied for improving UTD for off-policy reinforcement learning in a new paper

Install

$ pip install nGPT-pytorch

Usage

import torch
from nGPT_pytorch import nGPT

model = nGPT(
    num_tokens = 256,
    dim = 512,
    depth = 4,
    attn_norm_qk = True
)

x = torch.randint(0, 256, (2, 2048))

loss = model(x, return_loss = True)
loss.backward()

logits = model(x) # (2, 2048, 256)

Test

Enwik8

$ python train.py

Citations

@inproceedings{Loshchilov2024nGPTNT,
    title   = {nGPT: Normalized Transformer with Representation Learning on the Hypersphere},
    author  = {Ilya Loshchilov and Cheng-Ping Hsieh and Simeng Sun and Boris Ginsburg},
    year    = {2024},
    url     = {https://api.semanticscholar.org/CorpusID:273026160}
}

@article{Luo2017CosineNU,
    title     = {Cosine Normalization: Using Cosine Similarity Instead of Dot Product in Neural Networks},
    author    = {Chunjie Luo and Jianfeng Zhan and Lei Wang and Qiang Yang},
    journal   = {ArXiv},
    year      = {2017},
    volume    = {abs/1702.05870},
    url       = {https://api.semanticscholar.org/CorpusID:1505432}
}

@inproceedings{Zhou2024ValueRL,
    title   = {Value Residual Learning For Alleviating Attention Concentration In Transformers},
    author  = {Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan},
    year    = {2024},
    url     = {https://api.semanticscholar.org/CorpusID:273532030}
}

Metadata

Release files for nGPT-pytorch 0.2.16

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for nGPT-pytorch 0.2.16
File	Size	Uploaded
ngpt_pytorch-0.2.16.tar.gz	36.9 MB	Jun 3, 2025	Details

Built distribution (wheel)

Table of built distributions (wheels) for nGPT-pytorch 0.2.16
File	Interpreter	ABI	Platform	Reset
ngpt_pytorch-0.2.16-py3-none-any.whl 16.6 kB Jun 3, 2025	Python 3	none	any	Details

Total release size: 36.9 MB

Release files / ngpt_pytorch-0.2.16.tar.gz

Download URL	ngpt_pytorch-0.2.16.tar.gz
Size	36.9 MB
Tags	Source
SHA-256 checksum How to use checksums	`37ac18ee5b6de25448de67bc7a83cdaecc06d7ef32e013b1b2337894a20a02f7`
BLAKE2b-256 checksum How to use checksums	`e653232fd16e9d4bf2bda5a514714a82a1192e551de449230797e7a4e2ff5bd3`
Upload date	Jun 3, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.9.22`

Release files / ngpt_pytorch-0.2.16-py3-none-any.whl

Download URL	ngpt_pytorch-0.2.16-py3-none-any.whl
Size	16.6 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`01ca04c300994b427e9425acf4fa2343091a0d35627af7be57cbd1c2c6eeaf84`
BLAKE2b-256 checksum How to use checksums	`d310494fbda8fbfe497e4e653ee8b5156865818fa94be2c4ce79ced198e83d3e`
Upload date	Jun 3, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.9.22`