How to pack ternary numbers in 8-bit bytes

How to pack ternary numbers in 8-bit bytes with efficient SIMD-friendly unpacking

There are 3 possible values in a digit of a ternary number. 3 possible values, which could actually be anything. I've been recently nerd-sniped[^1] into trying to pack the ternary weights of BitNet b1.58 into something close to that theoretical ideal of log(3) / log(2) bits[^2] per ternary digit. I'll be calling a "ternary digit" a "trit", like a "binary digit" is called a "bit".

Block size

Since the goal of this is to allow fast parallel unpacking, blocks of trits can't be infinitely big. A small "block" size needs to be found, ideally one which is both efficient with information density and which is convenient on current hardware.

To find a good block size, we'll need to find a power of 3 for which the next power of 2 is very close.

It's very fortunate that 5 trits fit quite tight into 8 bits at 1.6 bits per trit. When compared to perfect packing, this is 99.06% efficient.

1.6 bits per trit

The basic idea with this packing scheme is simply to make a number out of the ternary digits.

def pack_number(digits: list[int], base: int) -> int:
    number = 0
    for digit in digits:
        assert digit < base
        number = number * base
        number += digit
    return number

Then, each 8-bit byte can store a 5-digit ternary number, effectively representing a fixed-point number between 0 and 1 inclusive, where the byte can store 243 possible values. The least significant base-3 digit is stored first.

def pack_ternary(digits: list[int]) -> bytearray:
    assert len(digits) % 5 == 0  # padding isn't handled here
    n_bytes = len(digits) // 5
    packed = bytearray()
    for i in range(n_bytes):
        b = 0
        for j in range(5):
            digit = digits[5*i + j]
            digit = max(-1, min(digit, 1))  # clamp between -1 and 1
            digit += 1  # from -1, 0, 1 to 0, 1, 2
            b *= 3
            b += digit
        b = ((b * 256) + (243 - 1)) // 243
        packed.append(b)
    return packed

The relevant interesting line is this one:

b = ((b * 256) + (243 - 1)) // 243

It does what is depicted in the diagram above, but multiplication is done first because these are integer operations. Doing a ceiling division here is necessary to cancel the off-by-one error from truncating when extracting digits later.

Unpacking

To unpack without using the modulo operator:

def unpack_trits(packed: bytes) -> list[int]:
    trits: list[int] = []
    for byte in packed:
        b = byte
        for i in range(5):
            b = b * 3
            trit = b >> 8
            trits.append(trit - 1)  # 0, 1, 2 => -1, 0, 1
            b = b & 0xFF
    return trits

Verification

To convince myself that this works, I wrote a C program checking that this really is lossless:

#include <stdio.h>
#include <stdint.h>
#include <string.h>

int main(void) {
    char s1[6] = {0};
    char s2[6] = {0};
    for (uint8_t i = 0; i < 243; i++) {
        // Turn the low 8 bits into a base-3 number (least significant first)
        uint8_t n = i;
        for (int j = 0; j < 5; j++) {
            s1[j] = (n % 3) + '0';
            n /= 3;
        }
        // Turn that number into a fixed-point number smaller than 1
        uint8_t q = (((uint16_t) i) * 256 + (243 - 1)) / 243;
        // This extracts the most significant digit first
        for (int j = 0; j < 5; j++) {
            uint16_t m = q * 3;
            s2[4 - j] = (m >> 8) + '0';
            q = m & 0xFF;
        }
        printf("%s, %s: %s\n", s1, s2, strcmp(s1, s2) == 0 ? "\033[1;32mPASS\033[0m" : "\033[1;31mFAIL\033[0m");
    }
    return 0;
}

Compile and run with:

$ gcc ternary-packing.c -o ternary-packing
$ ./ternary-packing

And I'm getting PASS for each of the 243 ternary numbers which fit in 8 bits.

And this is the technique used in the ternary types in llama.cpp for TriLMs and BitNet b1.58, for which the pull request is https://github.com/ggml-org/llama.cpp/pull/8151, with SIMD implementations for both AVX2 and ARM NEON.

[^1]: obviously referring to https://xkcd.com/356/, but the initial motivation actually started from this review comment I posted on the initial BitNet b1.58 pull-request for llama.cpp

[^2]: log(3) / log(2) is also known as 1.584962500721156.