От координат тензора к одному плоскому буферу

Сначала предскажите смещение, затем перейдите к плоскому буферу

Тензор позволяет интерпретировать набор скалярных значений как $n$-мерный объект. Начнём с формы:

[2, 2, 3]

У неё три оси, поэтому ранг тензора равен 3. Размеры осей означают, что вдоль оси 0 расположены два среза, в каждом срезе вдоль оси 1 есть две строки, а в каждой строке вдоль оси 2 — три позиции. Логически значения расположены так:

срез 0: [[10, 11, 12], [20, 21, 22]]
срез 1: [[30, 31, 32], [40, 41, 42]]

Эта реализация тензора не хранит вложенные объекты строк. Все значения находятся в одном плоском Vec<f64>:

[10, 11, 12, 20, 21, 22, 30, 31, 32, 40, 41, 42]

Прежде чем читать дальше, сделайте два предположения:

1. В какую позицию плоского буфера, отсчитываемую от нуля, попадёт координата [1,0,2] и какое значение там хранится?
2. Может ли координата [1,2,0] выбрать значение? Если нет, какая ось первой не пройдёт проверку?

При построчном хранении последняя координата меняется быстрее остальных. Шаг на единицу по оси 2 передвигает нас на один элемент плоского буфера. Полная строка содержит три значения, поэтому шаг по оси 1 передвигает нас на три элемента. Полный срез содержит две строки по три значения, поэтому шаг по оси 0 передвигает нас на шесть элементов. Эти расстояния и есть шаги по осям, измеряемые в элементах:

[6, 3, 1]

Теперь разложим вычисление для первой координаты:

Получились две величины разного рода. Смещение 8 — целочисленная позиция в плоском буфере. Значение 32 — число f64, которое хранится в этой позиции. Если назвать обе величины «индексом», потеряется операция, связывающая координаты с хранилищем.

Второе предположение отклоняется до любого обращения к буферу. Размер оси 1 равен 2, поэтому допустимы только индексы 0 и 1. Координата [1,2,0] содержит на этой оси индекс 2. Ошибка сообщает ось 1, индекс 2 и размер 2, не позволяя арифметике незаметно перейти в другую логическую строку.

Представьте движение по каждой оси отдельным слагаемым

Для тензора ранга $d$ координата преобразуется в смещение по правилу:

Координата содержит один индекс $i_k$, отсчитываемый от нуля, для каждой оси $k$. Шаг $s_k$ показывает, на сколько позиций изменится смещение в плоском буфере, если увеличить этот индекс на единицу, оставив остальные индексы неизменными. Построчные шаги выводятся справа налево. Последний шаг равен единице, а каждый предыдущий — произведению размеров всех осей справа от него с проверкой переполнения.

Для формы [2,2,3] справа от оси $0$ находится $2\cdot3=6$ элементов, справа от оси $1$ — $3$, а после оси $2$ ничего нет. Поэтому получаются шаги [6,3,1]. Они измеряются в элементах f64, а не в байтах.

Формула применима только к допустимой координате. Реализация сначала требует ровно $d$ индексов, а затем обходит оси слева направо и перед добавлением очередного слагаемого проверяет $0 \le i_k < \operatorname{shape}[k]$. Для [1,0,2] слагаемые равны 6, 0 и 2, поэтому смещение равно 8.

Две крайние формы уточняют соглашение о произведении. Форма [] имеет нулевой ранг и пустой список шагов, но представляет одно скалярное значение: пустая координата соответствует смещению 0. У формы [2,0,3] шаги [0,3,1] и нет значений. Нулевой размер оси делает тензор пустым, однако конструктор всё равно выводит и проверяет всю цепочку произведений для шагов.

Свяжите каждый символ с устройством тензора

Элементы формы и шаги отвечают на разные вопросы. shape[k] показывает, сколько позиций существует на оси $k$, а $s_k$ — на сколько элементов плоского буфера передвигает один шаг по этой оси. В нашем примере размер оси 0 равен 2, а шаг — 6; размер оси 1 тоже равен 2, но её шаг равен 3.

При $d=0$ в формуле нет слагаемых, поэтому сумма равна нулю. Так пустая координата скаляра получает смещение 0, а отдельный инвариант конструктора задаёт для скаляра одно хранящееся значение. Если размер какой-либо оси равен нулю, допустимого $i_k$ для неё не существует. Поэтому у пустого тензора нет полной допустимой координаты, хотя его шаги определены однозначно.

От биграммных счётчиков к обучаемым матрицам и тензорам внимания

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

Две ключевые работы — Bengio et al., A Neural Probabilistic Language Model и Vaswani et al., Attention Is All You Need.

Бенжио и соавторы описывают n-граммные модели как таблицы условных вероятностей с коротким контекстом, которые не используют сходство слов, а затем вводят нейронную языковую модель с матрицей $C$ размера «словарь × ширина признакового представления», содержащей обучаемые признаки слов, и матрицами параметров для предсказания следующего слова. Позднее Васвани и соавторы объединяют одновременно обрабатываемые запросы, ключи и значения в матрицы $Q$, $K$ и $V$ и с помощью обучаемых проекций параллельно вычисляют несколько голов внимания, после чего конкатенируют их выходы.

Модель Бенжио всё ещё предсказывает по фиксированному окну предшествующих слов. В небольшом примере этой главы матрица $C$ имеет форму $[|V|,m]=[5,3]$ и хранит по три обучаемых признака для каждого из пяти элементов словаря. Для контекста из двух слов скрытое преобразование $H$ имеет форму $[h,2m]=[4,6]$, а матрица $U$ формы $[|V|,h]=[5,4]$ преобразует четыре скрытые активации в пять оценок следующего слова. Эти небольшие формы лишь показывают роль матриц и не воспроизводят размеры экспериментов или полный набор параметров из статьи.

В примере Transformer каждый из тензоров $Q$, $K$ и $V$ для одной головы имеет форму $[\text{токены},\text{ширина головы}]=[2,3]$. Локальный стек тензоров $Q$ добавляет отдельную ось голов и имеет форму $[\text{головы},\text{токены},\text{ширина головы}]=[2,2,3]$. Статья обосновывает матричные вычисления и параллельные головы, но не требует именно такого порядка осей.

Явно заданные формы тензоров позволяют представить эмбеддинги, обучаемые веса, активации и промежуточные результаты внимания в постепенно расширяемой реализации декодера. Один непрерывный буфер с построчным хранением — локальное решение этого курса, а не требование какой-либо из двух статей.

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

Объедините проверку формы и координат в одной реализации

Tensor::from_vec принимает форму во владение как Vec<usize>, а буфер данных — как Vec<f64>. Он выводит построчные шаги справа налево, проверяя каждое умножение, а затем требует, чтобы длина буфера точно совпадала с вычисленным числом элементов. Для пустой формы расчёт начинается с единицы — нейтрального элемента умножения, — поэтому скаляру требуется ровно одно значение. Любой нулевой размер оси даёт нулевое число элементов.

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

pub(crate) fn checked_row_major_layout(
    shape: &[usize],
) -> Result<(Vec<usize>, usize), TensorError> {
    if shape.is_empty() {
        return Ok((Vec::new(), 1));
    }

    let mut strides = vec![1; shape.len()];
    for axis in (0..shape.len() - 1).rev() {
        strides[axis] = shape[axis + 1]
            .checked_mul(strides[axis + 1])
            .ok_or(TensorError::ShapeOverflow)?;
    }

    let element_count = shape[0]
        .checked_mul(strides[0])
        .ok_or(TensorError::ShapeOverflow)?;
    Ok((strides, element_count))
}

pub(crate) fn checked_offset(
    shape: &[usize],
    strides: &[usize],
    base_offset: usize,
    coordinate: &[usize],
) -> Result<usize, TensorError> {
    debug_assert_eq!(shape.len(), strides.len());

    if coordinate.len() != shape.len() {
        return Err(TensorError::RankMismatch {
            expected: shape.len(),
            actual: coordinate.len(),
        });
    }

    let mut offset = base_offset;
    for (axis, ((&index, &dimension), &stride)) in
        coordinate.iter().zip(shape).zip(strides).enumerate()
    {
        if index >= dimension {
            return Err(TensorError::IndexOutOfBounds {
                axis,
                index,
                dimension,
            });
        }
        let contribution = index
            .checked_mul(stride)
            .ok_or(TensorError::ShapeOverflow)?;
        offset = offset
            .checked_add(contribution)
            .ok_or(TensorError::ShapeOverflow)?;
    }
    Ok(offset)
}

impl Tensor {
    /// Builds a tensor after checking its row-major layout and buffer length.
    pub fn from_vec(shape: Vec<usize>, data: Vec<f64>) -> Result<Self, TensorError> {
        let (strides, expected) = checked_row_major_layout(&shape)?;
        let actual = data.len();
        if actual != expected {
            return Err(TensorError::DataLengthMismatch { expected, actual });
        }

        Ok(Self {
            data,
            shape,
            strides,
        })
    }

    /// Returns the number of logical axes.
    pub fn rank(&self) -> usize {
        self.shape.len()
    }

    /// Returns the extent of every axis.
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    /// Returns the row-major suffix-product stride for every axis.
    pub fn strides(&self) -> &[usize] {
        &self.strides
    }

    /// Returns the number of stored values.
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Reports whether the flat buffer stores no values.
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Borrows the contiguous value buffer.
    pub fn as_slice(&self) -> &[f64] {
        &self.data
    }

    /// Mutably borrows the contiguous value buffer without changing its length.
    pub fn as_mut_slice(&mut self) -> &mut [f64] {
        &mut self.data
    }

    /// Consumes the tensor and returns its contiguous value buffer.
    pub fn into_vec(self) -> Vec<f64> {
        self.data
    }

Второй фрагмент отвечает за проверку координаты. Метод offset сначала сравнивает длину координаты с рангом. Затем он обходит оси слева направо и сообщает первый индекс, который не меньше размера своей оси. Методы get и get_mut вызывают ту же операцию, поэтому чтение и изменение не могут использовать разные правила индексации.

    /// Maps one in-bounds coordinate to its row-major flat-buffer offset.
    pub fn offset(&self, coordinate: &[usize]) -> Result<usize, TensorError> {
        checked_offset(&self.shape, &self.strides, 0, coordinate)
    }

    /// Borrows the value at one checked coordinate.
    pub fn get(&self, coordinate: &[usize]) -> Result<&f64, TensorError> {
        let offset = self.offset(coordinate)?;
        Ok(&self.data[offset])
    }

    /// Mutably borrows the value at one checked coordinate.
    pub fn get_mut(&mut self, coordinate: &[usize]) -> Result<&mut f64, TensorError> {
        let offset = self.offset(coordinate)?;
        Ok(&mut self.data[offset])
    }

Публичные методы позволяют наблюдать инварианты, не давая форме и шагам разойтись. rank, shape, strides, len, is_empty и as_slice читают состояние тензора. as_mut_slice разрешает менять значения, но не размер буфера; into_vec забирает тензор во владение и возвращает принадлежавшие ему значения.

TensorError различает четыре вида ошибок. ShapeOverflow означает, что необходимое произведение не представимо в usize. DataLengthMismatch { expected, actual } означает, что произведение допустимо, но длина переданного буфера неверна. RankMismatch { expected, actual } проверяется раньше осей, а IndexOutOfBounds { axis, index, dimension } указывает первую недопустимую ось. Методы возвращают эти ошибки через Result; обычный неверный ввод не вызывает panic.

Первый фрагмент библиотеки показывает, как тензоры разных этапов развития языковых моделей представляются одним конкретным типом. Он создаёт небольшие тензоры параметров $C$, $H$ и $U$, тензоры активаций $Q$, $K$ и $V$ для одной головы и локальный стек тензоров $Q$ по оси голов. Все значения равны нулю, потому что здесь важны только форма, шаги и число элементов; сами операции моделей появятся в последующих главах.

/// Builds tiny shape-only stand-ins for parameters and activations from the LLM history.
pub fn llm_shape_history_fixture() -> Result<Vec<(&'static str, Tensor)>, TensorError> {
    Ok(vec![
        ("toy Bengio C", Tensor::from_vec(vec![5, 3], vec![0.0; 15])?),
        ("toy Bengio H", Tensor::from_vec(vec![4, 6], vec![0.0; 24])?),
        ("toy Bengio U", Tensor::from_vec(vec![5, 4], vec![0.0; 20])?),
        (
            "toy Transformer Q (one head)",
            Tensor::from_vec(vec![2, 3], vec![0.0; 6])?,
        ),
        (
            "toy Transformer K (one head)",
            Tensor::from_vec(vec![2, 3], vec![0.0; 6])?,
        ),
        (
            "toy Transformer V (one head)",
            Tensor::from_vec(vec![2, 3], vec![0.0; 6])?,
        ),
        (
            "toy Transformer Q head stack",
            Tensor::from_vec(vec![2, 2, 3], vec![0.0; 12])?,
        ),
    ])
}

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

/// Reconstructs the immutable, contiguous tensor used throughout the chapter.
pub fn frozen_tensor_fixture() -> Result<Tensor, TensorError> {
    Tensor::from_vec(FROZEN_SHAPE.to_vec(), FROZEN_VALUES.to_vec())
}

Следующий фрагмент main.rs сначала создаёт тензоры из исторического сравнения. Затем он вычисляет проверяемый доступ по координате [1,0,2], меняет значение по координате [0,1,1] на 99, создаёт скалярный и пустой тензоры и получает ошибки ранга, границ и переполнения. Код сразу после показанного фрагмента выводит эти значения. Последняя строка обозначает границу следующей главы: глава 9 добавит новую интерпретацию формы, шагов и базового смещения, сохранив тот же буфер.

    let llm_shapes = llm_shape_history_fixture()?;

    let mut tensor = frozen_tensor_fixture()?;
    let selected_offset = tensor.offset(&SELECTED_COORDINATE)?;
    let selected_value = *tensor.get(&SELECTED_COORDINATE)?;
    *tensor.get_mut(&[0, 1, 1])? = 99.0;

    let scalar = Tensor::from_vec(vec![], vec![7.0])?;
    let scalar_offset = scalar.offset(&[])?;
    let scalar_value = scalar.get(&[])?;
    let empty = Tensor::from_vec(vec![2, 0, 3], vec![])?;

    let rank_error = tensor.offset(&[1, 0]).unwrap_err();
    let bounds_error = tensor.offset(&INVALID_COORDINATE).unwrap_err();
    let overflow_error = Tensor::from_vec(vec![usize::MAX, 2], vec![]).unwrap_err();

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

    let tensor = frozen_tensor_fixture()?;
    let offset = tensor.offset(&SELECTED_COORDINATE)?;
    let value = tensor.get(&SELECTED_COORDINATE)?;
    let bounds_error = match tensor.offset(&INVALID_COORDINATE) {
        Err(error) => error,
        Ok(_) => {
            return Err(TensorError::IndexOutOfBounds {
                axis: 1,
                index: 2,
                dimension: 2,
            });
        }
    };
    let shape = usize_csv(tensor.shape());
    let strides = usize_csv(tensor.strides());
    let buffer = value_csv(tensor.as_slice());
    let coordinate = usize_csv(&SELECTED_COORDINATE);
    let slice0 = format!(
        "SLICE axis0=0 row0={} row1={}",
        row_csv(&tensor, 0, 0)?,
        row_csv(&tensor, 0, 1)?
    );
    let slice1 = format!(
        "SLICE axis0=1 row0={} row1={}",
        row_csv(&tensor, 1, 0)?,
        row_csv(&tensor, 1, 1)?
    );
    let terms = SELECTED_COORDINATE
        .iter()
        .zip(tensor.strides())
        .enumerate()
        .map(|(axis, (&index, &stride))| {
            format!(
                "TERM axis={axis} index={index} stride={stride} contribution={}",
                index * stride
            )
        })
        .collect::<Vec<_>>();

Трассировка сохраняет точную десятичную запись и порядок из вывода программы на Rust. Вместе с выводом Tensor она показывает, что срезы, смещение, выбранное значение и ошибка границ описывают одни и те же хранящиеся данные.

Проследите путь координаты через срезы, арифметику и хранилище

Изучите схему в четыре прохода:

1. Найдите срез 1, строку 0 и третью позицию в этой строке.
2. Сопоставьте координате [1,0,2] шаги [6,3,1] и проследите все три записанных вклада.
3. Найдите смещение 8 в плоском буфере и отделите его как позицию от значения 32.0.
4. Сравните допустимый индекс 0 на оси 1 с отклонённым индексом 2 в координате [1,2,0].

Обе таблицы срезов и плоский буфер содержат одни и те же двенадцать значений из трассировки Rust. Это разные представления только в объяснительном смысле: в главе 8 ещё нет API TensorView. Расчёт шагов показывает, почему переход по внешней оси пропускает шесть элементов, а по последней — один.

Ромб и двойная рамка отмечают смещение 8 независимо от цвета. Сначала сопоставьте этот маркер координате [1,0,2], а затем — значению 32.0 в плоском буфере.

Панель проверки границ — часть того же рассуждения. Она показывает ось 1, индекс 2 и размер 2: тензор отклоняет координату до вычисления смещения и чтения значения. Даже если для недопустимых индексов арифметика дала бы правдоподобный результат, доступ не стал бы допустимым.

Предскажите каждый результат работы с формой

Шаги выводятся справа налево, но при проверке доступа сначала проверяется ранг, а затем оси обходятся слева направо.

1. Выведите построчные шаги в элементах для формы [3,4,5].
2. Для формы [2,2,3], шагов [6,3,1] и зафиксированного буфера вычислите вклад каждой оси координаты [1,0,2]. Отдельно укажите смещение и значение.
3. Предскажите ошибку для координаты [1,2,0]. Какие ось, индекс и размер должны быть в ней указаны?
4. Для скалярной формы [] укажите ранг, шаги, число элементов, единственную допустимую координату и её смещение.
5. Для формы [2,0,3] выведите шаги и число элементов. Почему полная допустимая координата не существует?
6. Сравните два неудачных вызова конструктора: форма [2,2] с тремя значениями и форма [usize::MAX,2] с пустым вектором данных. Какая ошибка относится к каждому случаю?
7. Предскажите, какое смещение изменится, когда get_mut(&[0, 1, 1]) присвоит значение 99, и запишите получившийся плоский буфер.

Важно не перепутать: шаги не равны размерам осей. Размер показывает число допустимых позиций на оси, а шаг — перемещение по плоскому буферу при переходе на одну позицию по этой оси. Построчный порядок тоже является соглашением, а не неотъемлемым законом тензоров. Многомерные данные можно представить по столбцам или через непрямую неоднородную структуру с другим правилом преобразования координат в позиции хранилища.

Используйте один проверяемый числовой контейнер во всей модели

Теперь в курсе есть Tensor, которому принадлежит один непрерывный Vec<f64>. Форма задаёт его $n$-мерную интерпретацию, построчные шаги однозначно сопоставляют допустимым координатам смещения, а проверяемый доступ не даёт ошибкам ранга или границ превратиться в чтение несвязанных элементов буфера. В следующих главах на этом же инварианте хранения будут построены эмбеддинги, активации, веса, промежуточные результаты внимания, логиты, функции потерь и градиенты.

В главе 9 представления и преобразования осей позволят по-новому интерпретировать то же хранилище. Для этого могут появиться базовое смещение и нестандартные шаги, но буфер не должен неявно копироваться, а его значения — изменяться. Арифметика, расширение по осям, матричное умножение и дифференцирование относятся к последующим главам. Вклад главы 8 — устойчивое хранилище и проверяемая индексация, на которые они смогут опереться.