| """BitNetBase: 1-bit and ternary inference on hyperdimensional primitives. |
| |
| The observation this module is built on is that a 1-bit dot product and an |
| HDC similarity are *the same operation*. For sign vectors ``a, b`` in |
| {-1, +1}^n, packed one bit per weight with +1 -> 1 and -1 -> 0: |
| |
| disagree = popcount(a_bits ^ b_bits) # XOR is HDC's bind |
| dot = n - 2 * disagree # popcount is HDC's sim |
| |
| so ``sim(a, b) = dot(a, b) / n``. A BitNet layer's entire forward pass is |
| therefore XOR-and-bitcount over packed registers — no multiply, no float |
| accumulate — and :func:`dot_1bit_packed` verifies that identity against the |
| plain integer dot at runtime rather than asserting it in prose. |
| |
| Two things live here: |
| |
| ``BitLinear`` a ternary {-1, 0, +1} linear layer in the spirit of BitNet |
| b1.58: per-tensor absmean scaling with a dead zone, so small |
| weights collapse to 0 instead of costing compute |
| ``PrototypeNet`` a 1-bit classifier whose weights *are* hypervectors in an |
| :class:`~hdc.memory.HDCStorage` cell — training is bundling, |
| inference is a similarity argmax, and nothing is ever a |
| float array |
| |
| Standard library only; vectors are Python ints, so the kernels below are |
| bitwise ops over whole machine words. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import random |
| from typing import Sequence |
|
|
| from .hypervec import DEFAULT_DIM, bundle, hvc, mask, popcount, sim |
| from .memory import HDCStorage |
|
|
| __all__ = [ |
| "pack_signs", |
| "unpack_signs", |
| "dot_1bit", |
| "dot_1bit_packed", |
| "ternarize", |
| "quantize_absmax", |
| "ternary_dot", |
| "ternary_matvec", |
| "BitLinear", |
| "PrototypeNet", |
| "selfcheck", |
| ] |
|
|
|
|
| |
| |
| |
| def pack_signs(vec: Sequence[int]) -> int: |
| """Pack a {-1, +1} sequence into a bitmask, +1 -> 1, -1 -> 0. |
| |
| Bit ``i`` of the result is weight ``i``. Any non-positive entry counts as |
| -1, which keeps zero weights meaningful: a packed zero contributes -1 to |
| the dot product, so ternary matrices are packed with an explicit sparse |
| form (see :func:`ternary_dot`) rather than through this function. |
| """ |
| out = 0 |
| for i, v in enumerate(vec): |
| if v > 0: |
| out |= 1 << i |
| return out |
|
|
|
|
| def unpack_signs(bits: int, n: int) -> list[int]: |
| """Inverse of :func:`pack_signs`: ``n`` weights as +1/-1.""" |
| return [1 if (bits >> i) & 1 else -1 for i in range(n)] |
|
|
|
|
| def dot_1bit(a: Sequence[int], b: Sequence[int]) -> int: |
| """Plain integer dot product of two +/-1 sequences (the reference form).""" |
| if len(a) != len(b): |
| raise ValueError("length mismatch: %d vs %d" % (len(a), len(b))) |
| return sum(x * y for x, y in zip(a, b)) |
|
|
|
|
| def dot_1bit_packed(a_bits: int, b_bits: int, n: int) -> int: |
| """Bitwise dot product of two packed sign vectors. |
| |
| ``n - 2 * popcount(a ^ b)`` — one XOR and one bit-count for the whole |
| vector, independent of ``n`` in loop terms. This is the identity the |
| module exists to demonstrate: bind, then count disagreement. |
| """ |
| disagree = popcount((a_bits ^ b_bits) & mask(n)) |
| return n - 2 * disagree |
|
|
|
|
| |
| |
| |
| def ternarize(weights: Sequence[float], ratio: float = 0.7) -> list[int]: |
| """Quantize to {-1, 0, +1} with a per-tensor absmean scale. |
| |
| ``scale = mean(|w|)``, and weights inside ``ratio * scale`` of zero are |
| mapped to 0. With ``ratio = 1.0`` this is nearest-rounding of BitNet |
| b1.58's absmean scheme; lower ratios keep more weights non-zero (denser, |
| slightly more accurate), higher ratios sparsify the layer. |
| """ |
| items = [float(w) for w in weights] |
| if not items: |
| return [] |
| scale = sum(abs(w) for w in items) / len(items) |
| if scale == 0.0: |
| return [0] * len(items) |
| dead = ratio * scale |
| out: list[int] = [] |
| for w in items: |
| if w > dead: |
| out.append(1) |
| elif w < -dead: |
| out.append(-1) |
| else: |
| out.append(0) |
| return out |
|
|
|
|
| def quantize_absmax(x: Sequence[float], bits: int = 8) -> tuple[list[int], float]: |
| """Symmetric activation quantization to signed integers. |
| |
| Returns ``(integers, scale)`` such that ``x ~= integer * scale``. ``scale`` |
| is 1.0 for an all-zero input, which keeps the call total instead of |
| dividing by zero. |
| """ |
| if not 1 <= bits <= 16: |
| raise ValueError("bits must be in 1..16") |
| qmax = (1 << (bits - 1)) - 1 |
| peak = max((abs(float(v)) for v in x), default=0.0) |
| if peak == 0.0: |
| return [0] * len(x), 1.0 |
| scale = peak / qmax |
| out = [] |
| for v in x: |
| q = round(float(v) / scale) |
| out.append(max(-qmax - 1, min(qmax, q))) |
| return out, scale |
|
|
|
|
| def ternary_dot(w: Sequence[int], x: Sequence[int]) -> int: |
| """Integer dot product of a ternary weight row with an integer activation.""" |
| if len(w) != len(x): |
| raise ValueError("length mismatch: %d vs %d" % (len(w), len(x))) |
| total = 0 |
| for wi, xi in zip(w, x): |
| if wi == 1: |
| total += xi |
| elif wi == -1: |
| total -= xi |
| |
| return total |
|
|
|
|
| def ternary_matvec(matrix: Sequence[Sequence[int]], x: Sequence[int]) -> list[int]: |
| """Apply a ternary matrix to an integer vector.""" |
| return [ternary_dot(row, x) for row in matrix] |
|
|
|
|
| |
| |
| |
| class BitLinear: |
| """A ternary {-1, 0, +1} linear layer with additive-only accumulation. |
| |
| Weights are drawn once from a seeded normal distribution and ternarized at |
| construction, so two layers with the same shape and seed are identical. |
| ``forward`` accumulates with adds and subtracts only: a weight of -1 is a |
| subtraction and a weight of 0 is skipped, which is what makes 1-bit |
| inference cheap on hardware that has no multiplier. |
| |
| :meth:`scale` exposes the absmean scale discarded at ternarization time, so |
| callers that want a dequantized output can multiply it back in; ``forward`` |
| itself returns the raw integer accumulators. |
| """ |
|
|
| def __init__(self, in_features: int, out_features: int, seed: int = 0): |
| if in_features <= 0 or out_features <= 0: |
| raise ValueError("layer dimensions must be positive") |
| self.in_features = in_features |
| self.out_features = out_features |
| self.seed = seed |
| rng = random.Random(seed) |
| raw = [[rng.gauss(0.0, 1.0) for _ in range(in_features)] for _ in range(out_features)] |
| self.ternary_weights: list[list[int]] = [ternarize(row) for row in raw] |
| self._scales = [ |
| sum(abs(w) for w in row) / in_features for row in raw |
| ] |
|
|
| def scale(self, row: int) -> float: |
| """Absmean scale of output row ``row`` (multiply an output by this).""" |
| return self._scales[row] |
|
|
| def forward(self, x: Sequence[int]) -> list[int]: |
| """Integer accumulators for one input vector.""" |
| if len(x) != self.in_features: |
| raise ValueError( |
| "expected %d inputs, got %d" % (self.in_features, len(x)) |
| ) |
| return ternary_matvec(self.ternary_weights, x) |
|
|
| def density(self) -> float: |
| """Fraction of non-zero weights (measured, not assumed).""" |
| total = self.in_features * self.out_features |
| live = sum(1 for row in self.ternary_weights for w in row if w != 0) |
| return live / total if total else 0.0 |
|
|
|
|
| |
| |
| |
| class PrototypeNet: |
| """Nearest-prototype classifier with class weights stored as hypervectors. |
| |
| Training a class is bundling its samples into a single majority vector, so |
| the class weight is a superposition — the HDC centroid. Inference compares |
| a probe against each class cell by Hamming similarity and takes the argmax. |
| The weights are ordinary :class:`HDCStorage` cells, so they persist with |
| ``HDCStorage.save`` and are addressable by content like everything else. |
| """ |
|
|
| def __init__( |
| self, |
| storage: HDCStorage | None = None, |
| d: int = DEFAULT_DIM, |
| prefix: str = "proto/", |
| ): |
| self.d = d |
| self.prefix = prefix |
| self.storage = storage if storage is not None else HDCStorage(d=d) |
| self._samples: dict[str, list[int]] = {} |
|
|
| def _cell(self, label: str) -> str: |
| return self.prefix + label |
|
|
| def fit(self, label: str, sample: int) -> None: |
| """Add one hypervector sample to a class and rewrite its weight cell. |
| |
| The class weight is recomputed as the majority bundle of every sample |
| seen so far rather than folded in pairwise: pairwise bundling of two |
| vectors is AND-like and would forget the first sample quickly. |
| """ |
| samples = self._samples.setdefault(label, []) |
| samples.append(sample & mask(self.d)) |
| weight = bundle(samples, self.d) |
| self.storage.write(self._cell(label), weight, symbol=True) |
|
|
| def classes(self) -> list[str]: |
| """Class labels, in the order they were first fitted.""" |
| return list(self._samples) |
|
|
| def weights(self) -> list[tuple[str, int]]: |
| """``(label, weight vector)`` pairs as currently stored.""" |
| return [(label, self.storage.read(self._cell(label))) for label in self._samples] |
|
|
| def predict(self, sample: int) -> tuple[str, float]: |
| """Best matching class and its similarity; ``('', 0.0)`` when untrained.""" |
| best_label = "" |
| best = -2.0 |
| for label in self._samples: |
| s = sim(sample, self.storage.read(self._cell(label)), self.d) |
| if s > best: |
| best, best_label = s, label |
| if not best_label: |
| return "", 0.0 |
| return best_label, best |
|
|
| def accuracy(self, pairs: Sequence[tuple[str, int]]) -> float: |
| """Fraction of ``(expected_label, sample)`` pairs classified correctly.""" |
| if not pairs: |
| return 0.0 |
| hits = 0 |
| for label, sample in pairs: |
| got, _ = self.predict(sample) |
| if got == label: |
| hits += 1 |
| return hits / len(pairs) |
|
|
|
|
| |
| |
| |
| def _noisy(base: int, rate: float, rng: random.Random, d: int) -> int: |
| """Flip each bit of ``base`` with probability ``rate``.""" |
| v = base |
| for i in range(d): |
| if rng.random() < rate: |
| v ^= 1 << i |
| return v & mask(d) |
|
|
|
|
| def selfcheck() -> dict: |
| """Exercise the kernels and return measured values (nothing hardcoded). |
| |
| Covers the packed/plain dot-product equivalence, the sparsity a ternary |
| layer actually reaches, and end-to-end prototype classification accuracy |
| under 2% bit noise. |
| """ |
| rng = random.Random(20240501) |
|
|
| |
| trials = 25 |
| mismatches = 0 |
| for _ in range(trials): |
| n = rng.choice((16, 64, 256, 1000)) |
| signs_a = [rng.choice((-1, 1)) for _ in range(n)] |
| signs_b = [rng.choice((-1, 1)) for _ in range(n)] |
| plain = dot_1bit(signs_a, signs_b) |
| packed = dot_1bit_packed(pack_signs(signs_a), pack_signs(signs_b), n) |
| if plain != packed: |
| mismatches += 1 |
| xor_matches_dot = mismatches == 0 |
|
|
| |
| layer = BitLinear(128, 32, seed=7) |
| zeros = sum(1 for row in layer.ternary_weights for w in row if w == 0) |
| ternary_sparsity = zeros / (128 * 32) |
|
|
| |
| d = 2000 |
| net = PrototypeNet(d=d) |
| test_pairs: list[tuple[str, int]] = [] |
| for c in range(4): |
| base = hvc("selfcheck/class/%d" % c, d) |
| for _ in range(6): |
| net.fit("class%d" % c, _noisy(base, 0.02, rng, d)) |
| for _ in range(2): |
| test_pairs.append(("class%d" % c, _noisy(base, 0.02, rng, d))) |
|
|
| return { |
| "xor_matches_dot": xor_matches_dot, |
| "xor_trials": trials, |
| "ternary_sparsity": round(ternary_sparsity, 4), |
| "proto_density": round(layer.density(), 4), |
| "prototype_accuracy": round(net.accuracy(test_pairs), 4), |
| "prototype_classes": len(net.classes()), |
| "prototype_dim": d, |
| "d": DEFAULT_DIM, |
| } |
|
|
|
|
| if __name__ == "__main__": |
| print(json.dumps(selfcheck(), indent=2, sort_keys=True)) |
|
|
