В статье «Как исследователи нарушают привычные подходы в ИИ, исключая матричное умножение» упоминалось, в частности, что перспективным кажется хранение в нейросетевых матрицах лишь троичных значений: (-1, 0, 1), иначе говоря - тритов. Такие матрицы умножать друг на друга проще. И в моей статье я расскажу, как собственно, матрицы из тритов хранить и умножать.

Как известно, при умножении матриц, мы строку левой матрицы умножаем на столбец правой, и результат записываем в соответствующую ячейку результирующей матрицы. Чтобы было быстрее, мы правую матрицу предварительно транспонируем: тогда строку левой будем умножать на строку правой. Иначе говоря, мы скалярно перемножаем два вектора из тритов, причём оба занимают непрерывную область памяти.

Предположим, что процессор ориентирован на 32-битную арифметику. Тогда разобьём строку матрицы на векторы по 32 трита. Каждый из этих векторов будем хранить в виде двух 32-битных целых чисел, назовём их «плюс-вектором» и «минус-вектором». Трит с номером N равен разности бита с номером N плюс-вектора, и бита с номером N минус-вектора. При этом нулевое значение трита кодируется двумя способами: когда оба бита равны 0, или когда они оба равны 1.

трит

-1

0

0

1

плюс-вектор

0

0

1

1

минус-вектор

1

0

1

0

Такой способ хранения позволяет быстро подсчитать сумму всех тритов вектора. Для этого мы из суммы битов плюс-вектора вычитаем сумму битов минус-вектора. Различные алгоритмы нахождения суммы битов описаны в статье «Обстоятельно о подсчёте единичных битов», к тому же процессор может поддерживать инструкцию POPCNT, которая эту сумму подсчитывает.
Например, можно использовать следующий алгоритм из вышеупомянутой статьи:

// Количество единичных битов
unsigned __int32 popcnt(unsigned __int32 value)
{
    // Суммируем чётные и нечётные биты
    //  00  01  10  11  value
    //   0   0   1   1  (value >> 1) & 0x55555555
    //  --  --  --  --  --
    //  00  01  01  10  =
    value -= (value >> 1) & 0x55555555;

    // Повторяем уже для пар битов
    value = ((value >> 2) & 0x33333333) + (value & 0x33333333);

    // Умножение на 0x01010101 эквивалентно сумированию значений 4 байт числа,
    // при условии, что в младших байтах не будет переполнения
    // (а его в нашем случае не будет, так как там содержатся суммы битов).
    // Результат сложения будет в старшем байте произведения.
    return ((((value >> 4) + value) & 0x0F0F0F0F) * 0x01010101) >> 24;
}

Остаётся разобраться с потритовым умножением векторов. Для наглядности, будем считать трит структурой из двух битовых полей с именами p и m, причём p содержит значение из плюс-вектора, а m – из минус-вектора. Назовём «r» результат умножения тритов a и b. Тогда:
r.p = (a.p | b.m) & (a.m | b.p)
r.m = (a.p | b.p) & (a.m | b.m)
Итак, мы сначала по этим формулам потритово умножаем два троичных вектора, затем у результата этого умножения находим суммы битов плюс-вектора и минус-вектора, затем находим разность этих сумм. Это и будет результатом скалярного произведения двух троичных векторов друг на друга.

Проиллюстрируем этот алгоритм программой на языке C++. В этой программе мы создадим две троичные матрицы размерами 32x32, заполним их случайными значениями, и перемножим друг на друга двумя способами: классическим алгоритмом, и оптимизированным для троичных вычислений.

#include <stdlib.h>
#include <stdio.h>

// Неупакованные троичные матрицы,
// которые будем перемножать.
// Каждый элемент принимает значения (-1, 0, 1)
int A[32][32];
int B[32][32];

// Результат умножения неупакованных матриц
int C[32][32];

typedef unsigned __int32 u32;

// Вектор из 32 тритов
struct TritVector32
{
    u32 p; // плюс-вектор
    u32 m; // минус-вектор

    TritVector32() {p = 0; m = 0;}

    // Получить значение трита с указанным номером
    int getTrit(int index)
    {
        int mask = 1 << index;
        return ((p & mask) != 0) - ((m & mask) != 0);
    }
    void setTrit(int index, int trit)
    {
        int mask = 1 << index;
        p |= mask;
        p ^= mask & ((trit - 1) >> 1); // p ^= mask & -(trit <= 0);
        m |= mask;
        m ^= mask & ((~trit) >> 1);    // m ^= mask & -(trit >= 0);
    }
};

// Упакованные троичные матрицы,
// содержащие те же значения, что и неупакованные A и B
TritVector32 A3[32];
TritVector32 B3t[32]; // эта матрица транспонирована

// Результат умножения упакованных матриц
int C3[32][32];

// Возвращает случайное значение (-1, 0, 1)
int rand3()
{
    return (int)((u32) rand() * 3 / (RAND_MAX + 1)) - 1;
}

// Количество единичных битов
int popcnt(u32 value)
{
    // Для наглядности используется простой, но медленный алгоритм.
    // Существуют гораздо более быстрые.
    int result = 0;
    while(value)
    {
        result += value & 1;
        value >>= 1;
    }
    return result;
}

int main()
{
    // Заполняем матрицы сомножителей случайными значениями
    for(int i = 0; i < 32; i++)
        for(int j = 0; j < 32; j++)
        {
            A[i][j] = rand3();
            B[i][j] = rand3();
        }
    // Умножаем матрицы классическим способом
    for(int i = 0; i < 32; i++)
        for(int j = 0; j < 32; j++)
        {
            int c = 0;
            for(int k = 0; k < 32; k++)
                c += A[i][k] * B[k][j];
            C[i][j] = c;
        }

    // Заполняем упакованные матрицы
    for(int i = 0; i < 32; i++)
        for(int j = 0; j < 32; j++)
        {
            A3[i].setTrit(j, A[i][j]);
            B3t[j].setTrit(i, B[i][j]);
        }
    // Умножаем оптимизированным способом
    for(int i = 0; i < 32; i++)
    {
        TritVector32 *a = A3 + i;
        for(int j = 0; j < 32; j++)
        {
            TritVector32 *b = B3t + j;
            TritVector32 r;

            // Потритовое умножение
            r.p = (a->p | b->m) & (a->m | b->p);
            r.m = (a->p | b->p) & (a->m | b->m);

            C3[i][j] = popcnt(r.p) - popcnt(r.m);
        }
    }

    // Выводим результаты
    FILE *fp = fopen("classic.txt", "wt");
    fputs("A =\n", fp);
    for(int i = 0; i < 32; i++)
    {
        for(int j = 0; j < 32; j++)
            fprintf(fp, "%2d ", A[i][j]);
        fputs("\n", fp);
    }
    fputs("\nB =\n", fp);
    for(int i = 0; i < 32; i++)
    {
        for(int j = 0; j < 32; j++)
            fprintf(fp, "%2d ", B[i][j]);
        fputs("\n", fp);
    }
    fputs("\nC =\n", fp);
    for(int i = 0; i < 32; i++)
    {
        for(int j = 0; j < 32; j++)
            fprintf(fp, "%3d ", C[i][j]);
        fputs("\n", fp);
    }
    fclose(fp);

    fp = fopen("trit.txt", "wt");
    fputs("A =\n", fp);
    for(int i = 0; i < 32; i++)
    {
        for(int j = 0; j < 32; j++)
            fprintf(fp, "%2d ", A3[i].getTrit(j));
        fputs("\n", fp);
    }
    fputs("\nB =\n", fp);
    for(int i = 0; i < 32; i++)
    {
        for(int j = 0; j < 32; j++)
            fprintf(fp, "%2d ", B3t[j].getTrit(i));
        fputs("\n", fp);
    }
    fputs("\nC =\n", fp);
    for(int i = 0; i < 32; i++)
    {
        for(int j = 0; j < 32; j++)
            fprintf(fp, "%3d ", C3[i][j]);
        fputs("\n", fp);
    }
    fclose(fp);
    return 0;
}

Сравнение файлов результатов показывает, что они одинаковы. Всё работает.

Теперь возникает вопрос: а что, если у нас больше двух матриц? Ведь матрица произведения двух троичных матриц содержит числа, вообще говоря, не троичные. Как её умножить на троичную матрицу?
Но и здесь можно обойтись без умножения.
Пусть a – обычное целое число, которое может быть отрицательным. Пусть b – трит, имеющий два поля p и m, соответствующие плюс- и минус-векторам. Тогда произведение a на b можно записать так:
a•b = a•(b.p – b.m) = (a & -b.p) - (a & -b.m) =
= ((a - b.m) ^ -b.m) & -(b.m ^ b.p) =
= ((a - b.m) ^ -b.m) & ((-b.m) ^ -b.p) =
= ((a ^ -b.m) + b.m) & ((-b.m) ^ -b.p)
Можно выбрать любую из этих формул, или им аналогичных.

Комментарии (6)


  1. bodyawm
    12.11.2024 03:21

    Почему статья была без инвайта!?


    1. kinh Автор
      12.11.2024 03:21

      Не знал никого, кто бы мог мне инвайт дать. В своё время я интересовался троичной арифметикой, поэтому когда прочёл, что троичные матрицы предлагают использовать для нейросетей, то подумал, что, может, читателям будет интересно узнать про умножение таких матриц.


  1. MarkWatney
    12.11.2024 03:21

    При умножени двух троичных матрица, получаем обычную матрицу. Т.е. в результате сразу теряем нашу оптимизацию. И как сделать операцию нелинейной на тритах?

    Вообще не ясно как на таких матрицах можно нейросеть собрать


    1. MarkWatney
      12.11.2024 03:21

      Или это только оптимизации инференса (наподобие квантонизации)?


      1. kinh Автор
        12.11.2024 03:21

        Да, это оптимизация. Меньше размер матриц, больше скорость их умножения.


    1. kinh Автор
      12.11.2024 03:21

      как сделать операцию нелинейной на тритах?

      так как мы имеем дело с дискретными значениями, то их легко можно заменить с помощью табличной функции на любые другие. К примеру, умножаем троичную матрицу весов на какой-то входной троичный вектор. Результат - вектор из целых чисел. Используем каждое из этих чисел как номер элемента в массиве (номера могут быть отрицательными, поэтому предварительно можно сложить с какой-то константой bias), и подставляем вместо номера элемента его значение в этом массиве.
      На псевдокоде это можно примерно так записать:

      A: TritMatrix = [[1,0,0], [-1,1,-1], [0,1,-1]] // матрица весов
      x: TritVector = [1, 0, 1] // входной вектор
      z: TritVector = A * x // произведение матрицы на вектор
      f: [-3..3] int = [-2,-1,1,2,1,-1,-2] // табличная функция
      result: IntVector // целочисленный вектор результата
      for i: int = 0; i < 3; i++ {result[i] = f[z[i]]} // применение табличной функции
      return result // возврат результата