p‐kit Language Model Project

Working demos and more details are available at https://github.com/IBM/p-kit/tree/main/examples/llm We are also targetting running these demos on the Probana probilistic computer.

Goal

The goal of this project is to study whether a language model, and later potentially a larger LLM-like system, can be implemented over p-kit as an alternative to a conventional Transformer architecture.

The current implementation is a proof of concept based on a sparse stochastic p-bit recurrent reservoir. It has been evaluated first on synthetic character-language tasks and then on the full Tiny Shakespeare dataset. The main result is that the p-kit architecture is technically viable and exhibits different computational trade-offs from a Transformer.

The idea is to investigate whether p-kit can offer useful trade-offs in model simplicity, inference speed, trainable parameter count, hardware implementation, and energy efficiency.

Architecture

The current model follows this structure:

character/token
      ↓
sparse stochastic p-bit reservoir
      ↓
temporal reservoir history
      ↓
p-kit ReLU circuit
      ↓
reservoir + history + ReLU + token features
      ↓
linear readout
      ↓
softmax
      ↓
next-character probabilities

The recurrent reservoir is represented as a p-kit PCircuit with sparse pairwise couplings. Each token modifies the p-bit biases, while the previous reservoir state provides recurrent memory.

Several consecutive reservoir states can also be exposed directly to the readout. Experiments showed that this explicit temporal memory significantly improves language-model performance.

The reservoir and p-kit ReLU block are fixed. There is no backpropagation through the recurrent p-bit network. Training is concentrated in the final linear readout.

Preliminary results

The current full Tiny Shakespeare experiments show that a relatively small fixed stochastic p-bit recurrent network can learn meaningful character-language structure.

Compared with the small GPT reference model used in the experiments, the current p-kit implementation:

• uses substantially fewer trainable parameters;
• requires no gradient propagation through the recurrent p-bit network;
• currently achieves lower accuracy and higher perplexity than GPT;
• currently trains more slowly on CPU;
• achieves significantly faster CPU inference;
• still has substantial room for software and memory optimization.

The results should be considered indicative because both models are small. The main result is that the p-kit architecture is technically viable and exhibits different computational trade-offs from a Transformer.

Potential Advantages

Compared with a traditional Transformer language model, a p-kit-based architecture could potentially offer:

• Sparse computation — fewer active connections and lower memory traffic.
• Binary/stochastic state — potentially simpler and lower-precision hardware.
• Natural probabilistic sampling — stochastic behavior is built into the computation.
• Recurrent state by design — sequential context can be maintained without full attention over all previous tokens.
• Simpler recurrent training — no backpropagation through the fixed p-bit reservoir.
• Hardware flexibility — the same J, h, and p-bit model could target CPU, GPU, FPGA, MCU, or physical p-bit hardware.
• Potentially simpler specialized inference hardware — sparse stochastic elements may be easier to implement than a general Transformer accelerator.
• Energy-efficiency potential — especially if the model is mapped to dedicated stochastic or physical p-bit hardware.
• Reduced framework dependence — the core model does not fundamentally depend on PyTorch or TensorFlow.

The specialized-hardware aspect is particularly interesting: a future inference accelerator could potentially be built around arrays of simple stochastic p-bits with configurable sparse couplings instead of large attention and matrix-multiplication engines.

Limitations and Future Work

The current system remains a small character-level proof of concept. GPT still provides better predictive quality, and the scaling behavior of p-bit language models is not yet known.

The central research question is:

Can sparse stochastic p-bit networks become a useful alternative for implementing language models, especially when the model and specialized inference hardware are designed together?

The current project provides several demos that show that this direction is worth exploring.