16 · Версия материала 7

Выполните обратный проход по операциям, преобразующим ID токенов в значение функции потерь

Реализуйте VJP для матричных произведений, выбора строк эмбеддингов по повторяющимся ID, SiLU, log-softmax и средней функции потерь по токенам. Затем сравните каждое новое локальное правило с производными, оценёнными методом центральных разностей в выбранных координатах.

Предскажите путь повторяющегося токена

Начните с таблицы эмбеддингов из трёх строк, четырёх ID токенов, проекции в два класса и четырёх целевых классов:

E=\begin{bmatrix}2&2\\1&-1\\-1&1\end{bmatrix},\qquad z=\begin{bmatrix}1&1&1&2\end{bmatrix},\qquad \mathrm{targets}=\begin{bmatrix}0&0&0&1\end{bmatrix}

W=\begin{bmatrix}1&-1\\1&-1\end{bmatrix}

В формуле каждое вхождение токена обозначено парой $(b,t)$ : $b$ — индекс элемента пакета, в котором находится это вхождение, а $t$ — его позиция внутри элемента. В данном примере $B=1$ означает один элемент пакета, а $T=4$ — четыре позиции в нём. Пары перечислены в построчном порядке, поэтому плоский индекс равен $p=bT+t$ , а позиции $0,1,2,3$ соответствуют $(0,0),(0,1),(0,2),(0,3)$ .

Выбор строк материализует [1,-1] три раза и [-1,1] один раз. Умножение каждой выбранной строки на $W$ даёт значение проекции [0,0] до нелинейности. После SiLU этот вектор остаётся равным [0,0]. В этом компактном примере активированные значения непосредственно используются как логиты двух классов, поэтому log-softmax возвращает $[-\ln 2,-\ln 2]$ . Поэтому в каждой из четырёх позиций значение функции потерь равно $\ln 2$ , а среднее по позициям также равно $\ln 2$ .

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

Показанный log-softmax и совмещённое вычисление среднего NLL по индексам целевых классов — два отдельных вызова. У них общие входные логиты, но не результат прямого прохода. Каждый вызов сохраняет собственные данные, которые понадобятся его VJP.

До вычисления значений предскажите знаки градиентов. Чтобы уменьшить функцию потерь, логит целевого класса нужно увеличить, поэтому производная по нему отрицательна. Производная по логиту другого класса положительна. Здесь два логита в каждой позиции равны, поэтому вероятность каждого класса равна $1/2$ . После усреднения по четырём позициям модуль каждой компоненты градиента равен $1/8$ :

p\in\{0,1,2\},\ \mathrm{target}_p=0: \quad\begin{bmatrix}-\frac18&\frac18\end{bmatrix};\qquad p=3,\ \mathrm{target}_3=1: \quad\begin{bmatrix}\frac18&-\frac18\end{bmatrix}

Производная SiLU в нуле равна $1/2$ . Затем VJP проекции даёт три сопряжённые величины $[-1/8,-1/8]$ для трёх вхождений с ID 1 и одну $[1/8,1/8]$ для ID 2. Не усредняйте три вклада ID 1 ещё раз: множитель $1/4$ от средней функции потерь уже входит в каждый из них.

dE= \begin{bmatrix} 0 & 0 \\ -\frac{3}{8} & -\frac{3}{8} \\ \frac{1}{8} & \frac{1}{8} \end{bmatrix}, \qquad dW= \begin{bmatrix} -\frac{1}{4} & \frac{1}{4} \\ \frac{1}{4} & -\frac{1}{4} \end{bmatrix}

Сначала пример на Rust вручную вычисляет эту цепочку для заданных массивов. Затем тот же расчёт выражается через переиспользуемые локальные операции:

Вычислите вручную один обратный проход для заданных массивов, прежде чем вводить переиспользуемые VJP для отдельных операций rust/demos/ch16-model-autodiff-ops/src/lib.rs#handwritten-model-backward

/// Computes the compact token-operation chain with fixed arrays and handwritten rules.
///
/// This bounded reference calculation illustrates the model-specific backward
/// style that preceded a reusable operation vocabulary; it is not attributed
/// source code from any historical paper.
pub fn handwritten_model_backward() -> HandwrittenModelBaseline {
    let mut gathered = [0.0; 8];
    for (position, &token_id) in TOKEN_IDS.iter().enumerate() {
        for feature in 0..EMBEDDING_SHAPE[1] {
            gathered[position * 2 + feature] =
                EMBEDDING_VALUES[token_id * EMBEDDING_SHAPE[1] + feature];
        }
    }

    let mut projection_preactivations = [0.0; 8];
    for position in 0..TOKEN_IDS.len() {
        for class in 0..2 {
            for feature in 0..2 {
                projection_preactivations[position * 2 + class] +=
                    gathered[position * 2 + feature] * WEIGHT_VALUES[feature * 2 + class];
            }
        }
    }

    let mut activated = [0.0; 8];
    let mut log_probabilities = [0.0; 8];
    let mut loss = 0.0;
    let mut loss_input_gradient = [0.0; 8];
    let mut matmul_output_gradient = [0.0; 8];
    for position in 0..TOKEN_IDS.len() {
        for class in 0..2 {
            let offset = position * 2 + class;
            activated[offset] = projection_preactivations[offset]
                * stable_sigmoid(projection_preactivations[offset]);
        }
        let maximum = activated[position * 2].max(activated[position * 2 + 1]);
        let denominator = (activated[position * 2] - maximum).exp()
            + (activated[position * 2 + 1] - maximum).exp();
        for class in 0..2 {
            let offset = position * 2 + class;
            let probability = (activated[offset] - maximum).exp() / denominator;
            log_probabilities[offset] = probability.ln();
            let target_indicator = usize::from(TARGETS[position] == class) as f64;
            loss_input_gradient[offset] = (probability - target_indicator) / TOKEN_IDS.len() as f64;
            let sigmoid = stable_sigmoid(projection_preactivations[offset]);
            let silu_derivative =
                sigmoid * (1.0 + projection_preactivations[offset] * (1.0 - sigmoid));
            matmul_output_gradient[offset] = loss_input_gradient[offset] * silu_derivative;
        }
        loss -= log_probabilities[position * 2 + TARGETS[position]] / TOKEN_IDS.len() as f64;
    }

    let mut gathered_gradient = [0.0; 8];
    let mut weight_gradient = [0.0; 4];
    for position in 0..TOKEN_IDS.len() {
        for feature in 0..2 {
            for class in 0..2 {
                gathered_gradient[position * 2 + feature] += matmul_output_gradient
                    [position * 2 + class]
                    * WEIGHT_VALUES[feature * 2 + class];
                weight_gradient[feature * 2 + class] +=
                    gathered[position * 2 + feature] * matmul_output_gradient[position * 2 + class];
            }
        }
    }

    let mut embedding_gradient = [0.0; 6];
    for (position, &token_id) in TOKEN_IDS.iter().enumerate() {
        for feature in 0..2 {
            embedding_gradient[token_id * 2 + feature] += gathered_gradient[position * 2 + feature];
        }
    }

    HandwrittenModelBaseline {
        gathered,
        projection_preactivations,
        activated,
        log_probabilities,
        loss,
        loss_input_gradient,
        matmul_output_gradient,
        gathered_gradient,
        embedding_gradient,
        weight_gradient,
    }
}

Суммируйте вклады всех вхождений в общей строке таблицы эмбеддингов

Основная формула главы:

\frac{\partial L}{\partial E_{i,:}}=\sum_{(b,t):z_{b,t}=i}\frac{\partial L}{\partial X_{b,t,:}}

В прямом проходе операция выбора создаёт четыре независимые строки. Лента сохраняет ID каждого вхождения, поэтому при обратном проходе сопряжённая величина каждой строки возвращается в соответствующую строку исходной таблицы. Вклады вхождений с одинаковым ID суммируются в одной строке $\bar E$ : они не перезаписывают друг друга и не усредняются повторно.

Обозначьте функцию потерь, таблицу, ID токенов и сопряжённые величины

Символ	Смысл в вычислении
$L$	Скалярная средняя функция потерь по токенам.
$E$	Обучаемая таблица эмбеддингов формы $[V,d]$ .
$i$	Индекс строки таблицы $E$ , соответствующей одному элементу словаря.
$:$	Все координаты признаков указанной строки.
$b$	Индекс элемента пакета, в котором находится данное вхождение токена.
$t$	Позиция данного вхождения токена в элементе пакета.
$z_{b,t}$	ID токена в примере с индексом $b$ на позиции $t$ .
$X_{b,t,:}$	Выбранная строка признаков, используемая для этого вхождения.
$\partial L/\partial X_{b,t,:}$	Входящая сопряжённая величина этого вхождения.
$\sum_{(b,t):z_{b,t}=i}$	Суммировать все вхождения, ID которых равен $i$ .
$\partial L/\partial E_{i,:}$	Накопленная сопряжённая величина для всех признаков строки $i$ таблицы $E$ .

В главе 15 сопряжённая величина родителя обозначалась $\bar{x}$ . Здесь то же правило записывается как $\bar E_{i,:}\mathrel{+}=\bar X_{b,t,:}$ для каждого вхождения с ID $i$ . ID токенов не получают градиентов: они выбирают места назначения. Если вхождения хранятся плоско, формула $p=bT+t$ связывает каждую плоскую позицию с той же логической парой, по которой в сумме задано условие $z_{b,t}=i$ .

От обратного прохода в нейросетевой модели следующего слова к переиспользуемым VJP декодера

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

Этот ранний подход с обучаемыми векторами признаков слов описан в работе Bengio et al., A Neural Probabilistic Language Model. Bengio и соавторы строят нейросетевую модель следующего слова из обучаемых строк признаков слов, матричных уравнений, скрытого слоя tanh, нормированных вероятностей на выходе и явной фазы обратного прохода и обновления параметров модели.

Расчёт выше повторяет наглядный стиль обратного прохода, выписанного специально для одной модели. Это новый пример на Rust, а не код, взятый из статьи или приписанный ей. Для добавления другой операции пришлось бы изменить весь обратный расчёт.

Abadi и соавторы описывают графы тензорных операций. При автоматическом дифференцировании система прослеживает все пути от функции потерь к каждому параметру и суммирует их вклады в градиент, в том числе для выбранных строк эмбеддингов. В Transformer Vaswani и соавторов обучаемые эмбеддинги находятся на входе, а выходная проекция — на выходе модели; матричные проекции, softmax внимания и нелинейные подсети прямого распространения повторяются в стеке слоёв. Позднее Shazeer исследует Swish при $\beta=1$ — то есть SiLU — и варианты SwiGLU в подсетях прямого распространения Transformer.

Переход к переиспользуемым графам операций описан в работе Abadi et al., TensorFlow: A System for Large-Scale Machine Learning. Abadi и соавторы представляют операции вершинами графа, а тензоры — значениями рёбер, описывают автоматическое дифференцирование, которое суммирует все обратные пути к параметру, и показывают графы эмбеддингов с операцией Gather, градиенты которой обновляют выбранные строки.

Расположение операций в Transformer описано в работе Vaswani et al., Attention Is All You Need. Vaswani и соавторы строят Transformer из обучаемых эмбеддингов, обучаемых проекций запросов, ключей и значений, softmax внимания, сетей прямого распространения с двумя преобразованиями и ReLU, а также обучаемого выходного преобразования, за которым следует softmax.

Следующий этап развития сетей прямого распространения показан в работе Shazeer, GLU Variants Improve Transformer. Shazeer определяет Swish как произведение входа на сигмоиду от $\beta x$ ; при $\beta=1$ получается функция SiLU. В работе она используется в вариантах SwiGLU для подсетей прямого распространения Transformer; для исследованных вариантов с вентильным механизмом автор сообщает меньший логарифм перплексии на отложенной выборке, чем у рассмотренной базовой модели. В исходной работе о Transformer используется ReLU; она не подтверждает, что Vaswani и соавторы применяли SiLU.

В этой главе добавляются переиспользуемые локальные VJP для пакетных матричных произведений, выбора строк по повторяющимся ID, $\exp$ , $\log$ , SiLU, устойчивого log-softmax и совмещённого вычисления среднего NLL по индексам целевых классов. Log-softmax и совмещённое вычисление среднего NLL по индексам целевых классов сохраняют для своих VJP значения вероятностей, полученные при собственном вызове прямого прохода, и не нормализуют логиты повторно. Эти операции образуют набор локальных правил обратного прохода для последующих компонентов эмбеддингов, проекций, SwiGLU, внимания и функции потерь по токенам. При обычном инференсе выполняется только прямой проход.

Сохраняйте данные прямого прохода, необходимые каждому локальному VJP

Публичные методы добавляют операции в ленту из главы 15, не создавая отдельный граф. TensorValue::gather_rows — публичная точка входа, которая проверяет аргументы выбора строк. Сначала лента проверяет, что операнд с таблицей ещё доступен. Если операнд доступен, метод по порядку проверяет следующие условия:

таблица имеет ранг два;
для index_shape можно без переполнения вычислить раскладку в порядке по строкам и число логических позиций;
вычисленное число позиций равно indices.len(); и
каждый ID при обходе в плоском порядке меньше числа строк таблицы.

Если встречается ID вне диапазона строк, метод сообщает его первую плоскую позицию. Только после выполнения всех четырёх условий метод добавляет ширину таблицы к index_shape и проверяет полученную форму выхода. В рассматриваемом примере форма таблицы [3,2], значение index_shape=[4] и ID [1,1,1,2] дают форму выхода [4,2].

После успешной проверки закрытый тип RowGatherPlan, доступный только внутри крейта, владеет копиями ID и их логической формы, а также хранит форму исходной таблицы, вычисленную форму выхода и число элементов выхода. Закрытые поля не позволяют пользователю публичного API создать план в обход перечисленных проверок. Сам факт существования плана означает, что размеры и ID уже проверены; однако выделение памяти под выходной буфер всё ещё может завершиться ошибкой.

Один раз проверить запрос на выбор строк и создать план, хранящий проверенные ID и вычисленные формы rust/crates/llm-from-scratch/src/autograd/model_ops.rs#model-row-gather-plan

/// Owned row-gather facts established before materialization begins.
#[derive(Debug)]
pub(crate) struct RowGatherPlan {
    indices: Vec<usize>,
    index_shape: Vec<usize>,
    input_shape: [usize; 2],
    output_shape: Vec<usize>,
    output_len: usize,
}

impl RowGatherPlan {
    fn checked(
        table: &Tensor,
        indices: &[usize],
        index_shape: &[usize],
    ) -> Result<Self, TensorAutodiffError> {
        if table.rank() != 2 {
            return Err(ModelOpError::GatherTableRank { rank: table.rank() }.into());
        }
        let (_, expected) = checked_row_major_layout(index_shape)?;
        if indices.len() != expected {
            return Err(ModelOpError::GatherIndexCountMismatch {
                expected,
                actual: indices.len(),
            }
            .into());
        }
        let rows = table.shape()[0];
        for (position, &index) in indices.iter().enumerate() {
            if index >= rows {
                return Err(ModelOpError::GatherIndexOutOfBounds {
                    position,
                    index,
                    rows,
                }
                .into());
            }
        }

        Self::from_validated_indices(table, indices.to_vec(), index_shape.to_vec())
    }

    /// Seals indices whose shape, count, and bounds were established by a
    /// crate-owned caller for this exact rank-two table.
    pub(crate) fn from_validated_indices(
        table: &Tensor,
        indices: Vec<usize>,
        index_shape: Vec<usize>,
    ) -> Result<Self, TensorAutodiffError> {
        let input_shape = [table.shape()[0], table.shape()[1]];
        let width = input_shape[1];
        let mut output_shape = index_shape.clone();
        output_shape
            .try_reserve_exact(1)
            .map_err(|_| ModelOpError::OutputAllocationFailed {
                elements: indices.len().saturating_mul(width),
            })?;
        output_shape.push(width);
        let (_, output_len) = checked_row_major_layout(&output_shape)?;
        Ok(Self {
            indices,
            index_shape,
            input_shape,
            output_shape,
            output_len,
        })
    }

    fn into_saved_context(self) -> ModelSavedContext {
        let Self {
            indices,
            index_shape,
            input_shape,
            output_shape,
            ..
        } = self;
        ModelSavedContext::GatherRows {
            indices,
            index_shape,
            input_shape: input_shape.to_vec(),
            output_shape,
        }
    }
}

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

Создать проверенный план после проверки доступности операнда, скопировать строки и сохранить данные плана для обратного прохода rust/crates/llm-from-scratch/src/autograd/model_ops.rs#model-row-gather-operation

    /// Selects rows from a rank-two table into `index_shape + [width]`.
    ///
    /// IDs are integer selectors and are deliberately not tape operands.
    pub fn gather_rows(
        &self,
        indices: &[usize],
        index_shape: &[usize],
    ) -> Result<Self, TensorAutodiffError> {
        self.gather_rows_with_plan(|table| RowGatherPlan::checked(table, indices, index_shape))
    }

    /// Builds one row-gather plan after operand availability is established.
    pub(crate) fn gather_rows_with_plan(
        &self,
        build_plan: impl FnOnce(&Tensor) -> Result<RowGatherPlan, TensorAutodiffError>,
    ) -> Result<Self, TensorAutodiffError> {
        Self::model_operation(TensorOperation::GatherRows, [self], |primals| {
            let table = primals[0];
            let plan = build_plan(table)?;
            let value = gather_rows_forward(table, &plan)?;
            Ok((value, [plan.into_saved_context()]))
        })
    }

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

Добавить сопряжённую величину каждого вхождения в строку таблицы эмбеддингов, выбранную в прямом проходе rust/crates/llm-from-scratch/src/autograd/model_ops.rs#model-row-gather-vjp

        ModelSavedContext::GatherRows {
            indices,
            input_shape,
            output_shape,
            ..
        } => {
            debug_assert_eq!(upstream.shape(), output_shape);
            let width = input_shape[1];
            let mut table_gradient = zeros(input_shape)?;
            for (position, &index) in indices.iter().enumerate() {
                let source = position * width;
                let destination = index * width;
                for feature in 0..width {
                    table_gradient.as_mut_slice()[destination + feature] +=
                        upstream.as_slice()[source + feature];
                }
            }
            Ok(table_gradient)
        }

Остальные правила построены по тому же локальному принципу. VJP матричного умножения транспонируют последние две оси и суммируют вклады вдоль пакетных осей, расширенных при согласовании форм, чтобы восстановить исходную форму каждого родительского тензора. exp умножает входящую величину на сохранённый выход, log делит её на сохранённый положительный вход, а SiLU использует сохранённый вход и устойчиво вычисленное значение сигмоиды.

TensorValue::log_softmax выполняет один вызов прямого прохода с проверкой входных данных и получает два тензора: логарифмы вероятностей для вызывающего кода и обычные вероятности для VJP. После проверки входа один вызов обхода групп вычисляет для каждой группы один набор статистик и по этим данным формирует оба тензора. Операция сохраняет полученные значения вероятностей; при обратном проходе softmax повторно не вызывается и логиты заново не нормализуются.

Вернуть логарифмы вероятностей и сохранить вероятности из одного вызова прямого прохода с проверкой входных данных rust/crates/llm-from-scratch/src/autograd/model_ops.rs#model-log-softmax-saved-forward

    /// Applies stable log-softmax along one explicit class axis.
    pub fn log_softmax(&self, axis: usize) -> Result<Self, TensorAutodiffError> {
        Self::model_operation(TensorOperation::LogSoftmax, [self], |primals| {
            let input = primals[0];
            let forward = log_softmax_forward(&input.view(), axis, true)?;
            let probabilities = forward
                .probabilities
                .expect("the autodiff log-softmax forward requests saved probabilities");
            Ok((
                forward.value,
                [ModelSavedContext::LogSoftmax {
                    probabilities,
                    axis,
                    input_shape: input.shape().to_vec(),
                }],
            ))
        })
    }

TensorValue::indexed_mean_nll отдельно выполняет один вызов прямого прохода с проверкой входных данных для скалярного среднего NLL и вероятностей, необходимых VJP. Сначала лента подтверждает доступность операнда. После этого операция проверяет ось, число классов, число целевых индексов, непустой набор целей и каждый целевой индекс. Затем модуль вероятностей проверяет каждый логит на конечность и только после этого выделяет память под необязательный тензор вероятностей, сохраняя установленный порядок ошибок.

«Один вызов прямого прохода» не означает, что каждый логит читается один раз. Проверочный обход не вычисляет максимум, экспоненты, сумму или выход. Затем RowStats отдельно обходит классы для поиска максимума и для суммы сдвинутых экспонент. Обработчик группы использует эти данные, чтобы накопить скалярную функцию потерь, а код вывода ещё раз обходит классы и записывает тензор вероятностей для VJP. Сохранённый тензор содержит побитово те же значения f64, которые были сформированы при прямом проходе. Речь идёт о совпадении сохранённых чисел с плавающей запятой, а не о точных вещественных вероятностях. Все эти обходы относятся к одному вычислению нормализации, а не запускают нормализацию повторно.

Вернуть среднее NLL и сохранить вероятности из одного вызова прямого прохода с проверкой входных данных rust/crates/llm-from-scratch/src/autograd/model_ops.rs#model-indexed-nll-saved-forward

    /// Computes one stable rank-zero mean NLL from flat group-major targets.
    pub fn indexed_mean_nll(
        &self,
        axis: usize,
        targets: &[usize],
    ) -> Result<Self, TensorAutodiffError> {
        Self::model_operation(TensorOperation::IndexedMeanNll, [self], |primals| {
            let logits = primals[0];
            let forward = indexed_mean_nll_forward(&logits.view(), axis, targets, true)?;
            let probabilities = forward
                .probabilities
                .expect("the autodiff indexed-NLL forward requests saved probabilities");
            let value = Tensor::from_vec(Vec::new(), vec![forward.loss])?;
            Ok((
                value,
                [ModelSavedContext::IndexedMeanNll {
                    probabilities,
                    targets: targets.to_vec(),
                    axis,
                    input_shape: logits.shape().to_vec(),
                    groups: targets.len(),
                }],
            ))
        })
    }

Ветви log-softmax и среднего NLL по индексам целевых классов остаются двумя операциями. У каждой есть собственный результат и сохранённый контекст. Данные повторно используются между прямым проходом одной операции и её VJP, а не между двумя вызовами.

VJP для log-softmax вычитает из каждой входящей компоненты произведение сохранённой вероятности на сумму всех входящих компонент по оси классов. Пусть $g$ — группа, $c$ — класс, $Z_{g,c}$ — входной логит, $y_g$ — целевой класс, $G$ — число групп, $P_{g,c}$ — значение вероятности класса $c$ в группе $g$ , сохранённое при прямом проходе этой операции, а $\bar L$ — входящая скалярная сопряжённая величина. Тогда совмещённое вычисление среднего NLL по индексам целевых классов даёт сопряжённую величину логита по формуле

\bar Z_{g,c}=\frac{\bar L}{G}\left(P_{g,c}-\mathbf{1}[c=y_g]\right).

Для целевого класса из сохранённой вероятности в скобках вычитается единица. Затем все компоненты умножаются на общий множитель $\bar L/G$ : $\bar L$ передаёт входящую сопряжённую величину, а деление на $G$ учитывает усреднение по группам.

Полный пример строит выбор строк, матричное умножение, SiLU, отдельную ветвь log-softmax и совмещённое вычисление функции потерь из значений выше, а затем сопоставляет тензоры прямого и обратного проходов с расчётом на заданных массивах. Здесь вызывается backward_with_trace, потому что для сопоставления нужны сопряжённые величины выходов SiLU, матричного умножения и выбора строк. При обычном обучении используется метод без трассировки: backward, если подходит поведение сохранения графа по умолчанию, или backward_with_seed, если сохранение либо освобождение графа нужно выбрать явно. Все эти методы выполняют один и тот же обратный расчёт. Трассировка лишь записывает результаты локальных VJP; сами производные повторно не вычисляются.

Построить на общей ленте цепочку операций с повторяющимся токеном и сопоставить её с ручным расчётом для заданных массивов rust/demos/ch16-model-autodiff-ops/src/lib.rs#shared-model-vjp-fixture

/// Runs gather, matmul, SiLU, stable log-softmax, and combined token loss on one
/// repeated-token example, then checks the tape against the fixed reference.
pub fn frozen_model_example() -> Result<FrozenModelExample, TensorAutodiffError> {
    let embeddings = TensorValue::parameter(tensor(&EMBEDDING_SHAPE, &EMBEDDING_VALUES))?;
    let gathered = embeddings.gather_rows(&TOKEN_IDS, &TOKEN_SHAPE)?;
    let weights = TensorValue::parameter(tensor(&WEIGHT_SHAPE, &WEIGHT_VALUES))?;
    let projection_preactivations = gathered.matmul(&weights)?;
    let activated = projection_preactivations.silu()?;
    let log_probabilities = activated.log_softmax(CLASS_AXIS)?;
    let loss = activated.indexed_mean_nll(CLASS_AXIS, &TARGETS)?;
    let baseline = handwritten_model_backward();

    assert_close(gathered.value().as_slice(), &baseline.gathered, 1e-12);
    assert_close(
        projection_preactivations.value().as_slice(),
        &baseline.projection_preactivations,
        1e-12,
    );
    assert_close(activated.value().as_slice(), &baseline.activated, 1e-12);
    assert_close(
        log_probabilities.value().as_slice(),
        &baseline.log_probabilities,
        1e-12,
    );
    assert_close(loss.value().as_slice(), &[baseline.loss], 1e-12);

    let backward = loss.backward_with_trace()?;
    let loss_input_gradient = pass_adjoint(&backward, TensorOperation::Silu);
    let matmul_output_gradient = pass_adjoint(&backward, TensorOperation::MatMul);
    let gathered_gradient = pass_adjoint(&backward, TensorOperation::GatherRows);
    let embedding_gradient = embeddings
        .gradient_snapshot()
        .expect("embeddings are a parameter");
    let weight_gradient = weights
        .gradient_snapshot()
        .expect("weights are a parameter");
    assert_close(
        loss_input_gradient.as_slice(),
        &baseline.loss_input_gradient,
        1e-12,
    );
    assert_close(
        matmul_output_gradient.as_slice(),
        &baseline.matmul_output_gradient,
        1e-12,
    );
    assert_close(
        gathered_gradient.as_slice(),
        &baseline.gathered_gradient,
        1e-12,
    );
    assert_close(
        embedding_gradient.as_slice(),
        &baseline.embedding_gradient,
        1e-12,
    );
    assert_close(weight_gradient.as_slice(), &baseline.weight_gradient, 1e-12);

    Ok(FrozenModelExample {
        embeddings: embeddings.value_snapshot(),
        token_ids: TOKEN_IDS.to_vec(),
        weights: weights.value_snapshot(),
        targets: TARGETS.to_vec(),
        gathered: gathered.value_snapshot(),
        projection_preactivations: projection_preactivations.value_snapshot(),
        activated: activated.value_snapshot(),
        log_probabilities: log_probabilities.value_snapshot(),
        loss: loss.value_snapshot(),
        backward,
        loss_input_gradient,
        matmul_output_gradient,
        gathered_gradient,
        embedding_gradient,
        weight_gradient,
        baseline,
    })
}

Каждый новый VJP сравнивается с центральными разностями в выбранных координатах. Оба родителя матричного умножения проверяются отдельно; для каждой из остальных операций — выбора строк, exp, log, SiLU, log-softmax и среднего NLL по индексам — также выбираются отдельные координаты.

Сравнить каждый новый локальный VJP с центральными разностями в выбранных координатах rust/demos/ch16-model-autodiff-ops/src/lib.rs#model-vjp-gradchecks

/// Checks both matmul parents plus gather, exp, log, SiLU, log-softmax, and NLL.
pub fn model_vjp_gradchecks() -> Result<ModelVjpGradchecks, Box<dyn Error>> {
    let left = tensor(&[2, 2], &[0.4, -0.7, 1.2, 0.3]);
    let right = tensor(&[2, 2], &[0.2, -0.4, 0.9, 0.5]);
    let right_for_tape = right.clone();
    let right_for_objective = right.clone();
    let matmul_left = sampled_model_check(
        left.clone(),
        move |parameter| {
            let right = TensorValue::constant(right_for_tape)?;
            sum_to_scalar(parameter.matmul(&right)?)
        },
        move |candidate| {
            tensor_matmul(&candidate.view(), &right_for_objective.view())
                .expect("the worked matmul shapes agree")
                .as_slice()
                .iter()
                .sum()
        },
    )?;
    let left_for_tape = left.clone();
    let left_for_objective = left.clone();
    let matmul_right = sampled_model_check(
        right,
        move |parameter| {
            let left = TensorValue::constant(left_for_tape)?;
            sum_to_scalar(left.matmul(parameter)?)
        },
        move |candidate| {
            tensor_matmul(&left_for_objective.view(), &candidate.view())
                .expect("the worked matmul shapes agree")
                .as_slice()
                .iter()
                .sum()
        },
    )?;

    let gather_ids = [2, 1, 2];
    let gather = sampled_model_check(
        tensor(&[3, 2], &[0.2, -0.4, 0.7, 1.1, -0.3, 0.6]),
        |parameter| sum_to_scalar(parameter.gather_rows(&gather_ids, &[3])?),
        |candidate| raw_gather_sum(candidate, &gather_ids),
    )?;

    let exp = sampled_model_check(
        tensor(&[3], &[-0.8, 0.2, 1.1]),
        |parameter| sum_to_scalar(parameter.exp()?),
        |candidate| candidate.as_slice().iter().map(|value| value.exp()).sum(),
    )?;
    let log = sampled_model_check(
        tensor(&[3], &[0.4, 1.1, 2.3]),
        |parameter| sum_to_scalar(parameter.log()?),
        |candidate| candidate.as_slice().iter().map(|value| value.ln()).sum(),
    )?;
    let silu = sampled_model_check(
        tensor(&[3], &[-0.8, 0.2, 1.1]),
        |parameter| sum_to_scalar(parameter.silu()?),
        |candidate| {
            candidate
                .as_slice()
                .iter()
                .map(|&value| value * stable_sigmoid(value))
                .sum()
        },
    )?;

    let log_softmax_weights = tensor(&[2, 3], &[0.2, -0.5, 0.7, 1.1, -0.4, 0.3]);
    let weights_for_tape = log_softmax_weights.clone();
    let weights_for_objective = log_softmax_weights.clone();
    let log_softmax = sampled_model_check(
        tensor(&[2, 3], &[0.7, -0.4, 1.1, -0.2, 0.3, 0.8]),
        move |parameter| {
            let weights = TensorValue::constant(weights_for_tape)?;
            sum_to_scalar(parameter.log_softmax(1)?.mul(&weights)?)
        },
        move |candidate| {
            let output = tensor_log_softmax(&candidate.view(), 1)
                .expect("the worked probability axis is valid");
            weighted_sum(output.as_slice(), weights_for_objective.as_slice())
        },
    )?;
    let indexed_mean_nll = sampled_model_check(
        tensor(&[2, 3], &[0.7, -0.4, 1.1, -0.2, 0.3, 0.8]),
        |parameter| parameter.indexed_mean_nll(1, &[2, 0]),
        |candidate| {
            tensor_indexed_mean_nll(&candidate.view(), 1, &[2, 0])
                .expect("the worked targets are valid")
        },
    )?;

    let checks = vec![
        NamedModelGradcheck {
            operation: "matmul-left",
            report: matmul_left,
        },
        NamedModelGradcheck {
            operation: "matmul-right",
            report: matmul_right,
        },
        NamedModelGradcheck {
            operation: "gather_rows",
            report: gather,
        },
        NamedModelGradcheck {
            operation: "exp",
            report: exp,
        },
        NamedModelGradcheck {
            operation: "log",
            report: log,
        },
        NamedModelGradcheck {
            operation: "silu",
            report: silu,
        },
        NamedModelGradcheck {
            operation: "log_softmax",
            report: log_softmax,
        },
        NamedModelGradcheck {
            operation: "indexed_mean_nll",
            report: indexed_mean_nll,
        },
    ];
    let passed = checks.iter().all(|check| check.report.passed);
    Ok(ModelVjpGradchecks { checks, passed })
}

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

Отклонить недопустимые ID, целевые классы, пустые наборы целей и переполнение, сохранив записанные градиенты rust/demos/ch16-model-autodiff-ops/src/lib.rs#model-op-errors-example

pub fn model_error_example() -> Result<ModelErrorExample, TensorAutodiffError> {
    let table = TensorValue::parameter(tensor(&[2, 2], &[1.0, 2.0, 3.0, 4.0]))?;
    let logits = TensorValue::parameter(tensor(&[2, 2], &[1.0, -1.0, 0.5, -0.5]))?;
    let empty_logits = TensorValue::parameter(tensor(&[0, 2], &[]))?;
    let overflow = TensorValue::parameter(tensor(&[], &[f64::MAX]))?;

    sum_to_scalar(table.clone())?.backward()?;
    sum_to_scalar(logits.clone())?.backward()?;
    overflow.backward()?;

    let invalid_id_before = tensor_bits(&table.gradient().expect("table gradient exists"));
    let invalid_id = table
        .gather_rows(&[2], &[1])
        .expect_err("row two is outside a two-row table");
    let invalid_id_gradient_unchanged =
        invalid_id_before == tensor_bits(&table.gradient().expect("table gradient remains"));

    let invalid_target_before = tensor_bits(&logits.gradient().expect("logit gradient exists"));
    let invalid_target = logits
        .indexed_mean_nll(1, &[0, 2])
        .expect_err("class two is outside a two-class row");
    let invalid_target_gradient_unchanged =
        invalid_target_before == tensor_bits(&logits.gradient().expect("logit gradient remains"));

    let empty_targets_before =
        tensor_bits(&empty_logits.gradient().expect("empty gradient exists"));
    let empty_targets = empty_logits
        .indexed_mean_nll(1, &[])
        .expect_err("a mean over no targets is undefined");
    let empty_targets_gradient_unchanged = empty_targets_before
        == tensor_bits(&empty_logits.gradient().expect("empty gradient remains"));

    let exp_overflow_before = tensor_bits(&overflow.gradient().expect("scalar gradient exists"));
    let exp_overflow = overflow
        .exp()
        .expect_err("the finite-forward invariant rejects positive infinity");
    let exp_overflow_gradient_unchanged =
        exp_overflow_before == tensor_bits(&overflow.gradient().expect("scalar gradient remains"));
    let gradients_unchanged = invalid_id_gradient_unchanged
        && invalid_target_gradient_unchanged
        && empty_targets_gradient_unchanged
        && exp_overflow_gradient_unchanged;

    Ok(ModelErrorExample {
        invalid_id,
        invalid_target,
        empty_targets,
        exp_overflow,
        invalid_id_gradient_unchanged,
        invalid_target_gradient_unchanged,
        empty_targets_gradient_unchanged,
        exp_overflow_gradient_unchanged,
        gradients_unchanged,
    })
}

Программа из следующей вставки выводит рассмотренные значения, три скалярные производные, результат проверки центральными разностями и общий результат проверки сохранности градиентов при ошибках:

Выполнить полную цепочку операций и вывести результаты прямого и обратного проходов, численных проверок и ошибок rust/demos/ch16-model-autodiff-ops/src/main.rs#learner-model-vjp-output

    let example = frozen_model_example()?;
    let probes = scalar_probes()?;
    let gradchecks = model_vjp_gradchecks()?;
    let errors = model_error_example()?;

    println!("embeddings: {}", tensor_text(&example.embeddings));
    println!("token IDs: {:?}", example.token_ids);
    println!("gather rows: {}", tensor_text(&example.gathered));
    println!("projection weights: {}", tensor_text(&example.weights));
    println!(
        "projection preactivations: {}",
        tensor_text(&example.projection_preactivations)
    );
    println!("SiLU: {}", tensor_text(&example.activated));
    println!(
        "log-softmax axis=1: {}",
        tensor_text(&example.log_probabilities)
    );
    println!("targets: {:?}", example.targets);
    println!("indexed mean NLL: {}", tensor_text(&example.loss));
    println!(
        "target-logit gradient: {}",
        tensor_text(&example.loss_input_gradient)
    );
    println!(
        "through SiLU: {}",
        tensor_text(&example.matmul_output_gradient)
    );
    println!(
        "matmul left gradient: {}",
        tensor_text(&example.gathered_gradient)
    );
    println!(
        "embedding scatter-add: {}",
        tensor_text(&example.embedding_gradient)
    );
    println!(
        "matmul right gradient: {}",
        tensor_text(&example.weight_gradient)
    );
    println!(
        "scalar probes: exp(0)->({}, {}) | log(1)->({}, {}) | silu(0)->({}, {})",
        fixed(probes[0].output),
        fixed(probes[0].gradient),
        fixed(probes[1].output),
        fixed(probes[1].gradient),
        fixed(probes[2].output),
        fixed(probes[2].gradient),
    );
    println!(
        "gradcheck: {}; pass={}",
        gradchecks
            .checks
            .iter()
            .map(|check| check.operation)
            .collect::<Vec<_>>()
            .join(" | "),
        gradchecks.passed
    );
    println!(
        "typed errors: invalid-id | invalid-target | empty-targets | exp-overflow; gradients unchanged={}",
        errors.gradients_unchanged
    );
    println!("chapter 17 handoff: initialize trainable values reproducibly");

embeddings: shape=3x2 values=2.000000000000,2.000000000000,1.000000000000,-1.000000000000,-1.000000000000,1.000000000000
token IDs: [1, 1, 1, 2]
gather rows: shape=4x2 values=1.000000000000,-1.000000000000,1.000000000000,-1.000000000000,1.000000000000,-1.000000000000,-1.000000000000,1.000000000000
projection weights: shape=2x2 values=1.000000000000,-1.000000000000,1.000000000000,-1.000000000000
projection preactivations: shape=4x2 values=0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000
SiLU: shape=4x2 values=0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000
log-softmax axis=1: shape=4x2 values=-0.693147180560,-0.693147180560,-0.693147180560,-0.693147180560,-0.693147180560,-0.693147180560,-0.693147180560,-0.693147180560
targets: [0, 0, 0, 1]
indexed mean NLL: shape=scalar values=0.693147180560
target-logit gradient: shape=4x2 values=-0.125000000000,0.125000000000,-0.125000000000,0.125000000000,-0.125000000000,0.125000000000,0.125000000000,-0.125000000000
through SiLU: shape=4x2 values=-0.062500000000,0.062500000000,-0.062500000000,0.062500000000,-0.062500000000,0.062500000000,0.062500000000,-0.062500000000
matmul left gradient: shape=4x2 values=-0.125000000000,-0.125000000000,-0.125000000000,-0.125000000000,-0.125000000000,-0.125000000000,0.125000000000,0.125000000000
embedding scatter-add: shape=3x2 values=0.000000000000,0.000000000000,-0.375000000000,-0.375000000000,0.125000000000,0.125000000000
matmul right gradient: shape=2x2 values=-0.250000000000,0.250000000000,0.250000000000,-0.250000000000
scalar probes: exp(0)->(1.000000000000, 1.000000000000) | log(1)->(0.000000000000, 1.000000000000) | silu(0)->(0.000000000000, 0.500000000000)
gradcheck: matmul-left | matmul-right | gather_rows | exp | log | silu | log_softmax | indexed_mean_nll; pass=true
typed errors: invalid-id | invalid-target | empty-targets | exp-overflow; gradients unchanged=true
chapter 17 handoff: initialize trainable values reproducibly

Проследите четыре градиентных вклада до трёх строк параметра

Схема отделяет компактную цепочку прямого прохода от накопления при обратном проходе. После SiLU граф явно разветвляется на два отдельных вызова над одними и теми же логитами: log-softmax возвращает логарифмы вероятностей, а совмещённое вычисление среднего NLL — скалярную функцию потерь. Каждая операция сохраняет значения вероятностей, полученные при собственном вызове прямого прохода. В строках целей видны отрицательная и положительная производные, их нулевая сумма по классам и точные формы сопряжённых величин, передаваемых назад через SiLU и к обоим операндам матричного умножения.

В блоках строк назначения находятся вклады соответствующих вхождений. Строка 1 содержит позиции 0, 1 и 2 до их суммирования, строка 2 — позицию 3, а строка 0 явно отмечена как неиспользованная. Так видно, почему после обратного прохода градиент родительской таблицы содержит три строки, хотя операция выбора материализовала четыре выходные строки.

Проследите путь повторяющегося ID токена до функции потерь и обратно

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

ID токенов: [1, 1, 1, 2]
Целевые классы: [0, 0, 0, 1]
Средняя функция потерь по токенам: 0.693147180560
Повторяющийся ID: $i=1,\;n=3$

Проследите прямой проход по компактной цепочке операций

После SiLU граф разветвляется на два отдельных вызова над одними и теми же логитами. Log-softmax возвращает логарифмы вероятностей, а совмещённая операция вычисляет по логитам и целевым классам скалярное среднее NLL. Вызовы не используют общий результат прямого прохода. Эта схема предназначена для изучения операций, а не устройства итогового декодера.

Выбор строк

Шаг прямого прохода
0

Входы
таблица эмбеддингов + ID токенов

Форма дифференцируемого тензора

$[3,2]$

Форма выхода
$[4,2]$

Значения
[1.000000000000, -1.000000000000, 1.000000000000, -1.000000000000, 1.000000000000, -1.000000000000, -1.000000000000, 1.000000000000]
Матричное умножение

Шаг прямого прохода
1

Входы
выбранные строки признаков + веса проекции

Форма дифференцируемого тензора

$[4,2]$ $[2,2]$

Форма выхода
$[4,2]$

Значения
[0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000]
SiLU

Шаг прямого прохода
2

Входы
выходы проекции до нелинейности

Форма дифференцируемого тензора

$[4,2]$

Форма выхода
$[4,2]$

Значения
[0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000]
Log-softmax

Шаг прямого прохода
3

Входы
выходы SiLU, используемые как логиты функции потерь

Форма дифференцируемого тензора

$[4,2]$

Форма выхода
$[4,2]$

Значения
[-0.693147180560, -0.693147180560, -0.693147180560, -0.693147180560, -0.693147180560, -0.693147180560, -0.693147180560, -0.693147180560]
Среднее NLL по индексам целевых классов

Шаг прямого прохода
4

Входы
выходы SiLU, используемые как логиты функции потерь + целевые классы

Форма дифференцируемого тензора

$[4,2]$

Форма выхода
$[]$

Значения
[0.693147180560]

Проведите обратный проход через выбор целевого класса и проекцию

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

Градиенты функции потерь по логитам в четырёх плоских позициях
Плоская позиция	ID токена	Целевой класс	Градиент	Знак градиента для целевого класса	Знак градиента для другого класса
0	1	0	`[-0.125000000000, 0.125000000000]`	отрицательный — выбранный целевой класс	положительный
1	1	0	`[-0.125000000000, 0.125000000000]`	отрицательный — выбранный целевой класс	положительный
2	1	0	`[-0.125000000000, 0.125000000000]`	отрицательный — выбранный целевой класс	положительный
3	2	1	`[0.125000000000, -0.125000000000]`	отрицательный — выбранный целевой класс	положительный

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

SiLU

Операнд: единственный вход
Родительский тензор: выходы проекции до нелинейности
Форма градиента по родительскому тензору: $[4,2]$
Градиент: [-0.062500000000, 0.062500000000, -0.062500000000, 0.062500000000, -0.062500000000, 0.062500000000, 0.062500000000, -0.062500000000]

Матричное умножение

Операнд: левая матрица
Родительский тензор: выбранные строки признаков
Форма градиента по родительскому тензору: $[4,2]$
Градиент: [-0.125000000000, -0.125000000000, -0.125000000000, -0.125000000000, -0.125000000000, -0.125000000000, 0.125000000000, 0.125000000000]

Матричное умножение

Операнд: правая матрица
Родительский тензор: веса проекции
Форма градиента по родительскому тензору: $[2,2]$
Градиент: [-0.250000000000, 0.250000000000, 0.250000000000, -0.250000000000]

Добавьте вклад каждого вхождения в строку назначения

У каждой выбранной выходной строки есть собственная копия значений. При обратном проходе для каждого вхождения сохраняется отдельная сопряжённая величина, а VJP выбора строк добавляет её в строку родительской таблицы, номер которой равен ID токена.

неиспользованная строка

Строка назначения: 0
Плоские позиции: неиспользованная строка
Число вхождений: 0
Градиент: [0.000000000000, 0.000000000000]

строка с суммой вкладов

Строка назначения: 1
Плоские позиции: 0, 1, 2
Число вхождений: 3
Градиент: [-0.375000000000, -0.375000000000]

вхождение повторяющегося ID

Плоская позиция
0

Вклад вхождения
[-0.125000000000, -0.125000000000]
вхождение повторяющегося ID

Плоская позиция
1

Вклад вхождения
[-0.125000000000, -0.125000000000]
вхождение повторяющегося ID

Плоская позиция
2

Вклад вхождения
[-0.125000000000, -0.125000000000]

строка с одним вкладом

Строка назначения: 2
Плоские позиции: 3
Число вхождений: 1
Градиент: [0.125000000000, 0.125000000000]

единственное вхождение ID

Плоская позиция
3

Вклад вхождения
[0.125000000000, 0.125000000000]

Предскажите результат до запуска Rust

Сопоставьте каждую плоскую позицию от 0 до 3 с парой $(b,t)$ при $B=1$ и $T=4$ .
Запишите четыре выбранные строки, строки проекции до нелинейности, строки логитов, подаваемых в функцию потерь, строки логарифмов вероятностей и скалярную среднюю функцию потерь.
Объясните, почему градиент по логиту целевого класса отрицателен, по логиту другого класса — положителен, а сумма компонент градиента по двум классам в каждой строке равна нулю.
Примените производную SiLU в нуле. Какая сопряжённая величина поступает к VJP матричного умножения для каждой позиции?
Предскажите формы градиентов обоих родителей матричного умножения и вычислите $dW$ .
Перечислите три вклада в строку 1 таблицы эмбеддингов до их суммирования. Какие значения получат строки 0 и 2?
Предскажите значение и локальную производную для $\exp(0)$ , $\ln(1)$ и $\operatorname{SiLU}(0)$ .
Объясните, почему при логитах около $\pm1000$ нужна нормализация с вычитанием максимума и почему каждая дифференцируемая операция с вероятностями сохраняет значения, полученные при собственном вызове прямого прохода после проверки входных данных. Означает ли «один вызов», что каждый логит читается один раз? В каком смысле сохранённые значения f64 совпадают и используют ли две ветви общий результат?
Перечислите свойства запроса на выбор строк, которые проверяет публичная точка входа. Какие из этих проверок не повторяются после создания проверенного плана? На каком этапе обнаруживаются неверный целевой класс и пустой набор целей?
Распространённая ошибка: нужно ли делить строку 1 на три после сложения трёх вкладов от вхождений с ID 1?

Проверьте десять предсказаний об операциях модели

Получается $0\mapsto(0,0)$ , $1\mapsto(0,1)$ , $2\mapsto(0,2)$ и $3\mapsto(0,3)$ .
Выбор строк даёт [1,-1] три раза и [-1,1] один раз. Каждая строка проекции до нелинейности, выхода SiLU и логитов функции потерь равна [0,0]; каждая строка логарифмов вероятностей равна $[-\ln 2,-\ln 2]$ ; средняя функция потерь равна $\ln 2$ .
Увеличение логита выбранного целевого класса уменьшает NLL, поэтому производная по нему отрицательна. Увеличение логита другого класса повышает его вероятность за счёт целевого класса, поэтому производная положительна. Сумма вероятностей softmax равна единице, как и сумма компонент индикаторного вектора целевого класса, поэтому сумма компонент градиента по оси классов равна нулю.
Производная SiLU в нуле равна $1/2$ . Первые три строки становятся $[-1/16,1/16]$ , последняя — $[1/16,-1/16]$ .
Градиент левого родителя имеет форму выбранных строк [4,2], а градиент правого родителя — форму весов [2,2]. Градиент весов равен $dW=[[-1/4,1/4],[1/4,-1/4]]$ .
Позиции 0, 1 и 2 вносят по $[-1/8,-1/8]$ в строку 1, поэтому она получает $[-3/8,-3/8]$ . Строка 2 получает $[1/8,1/8]$ , а неиспользованная строка 0 — [0,0].
$\exp(0)=1$ с производной $1$ ; $\ln(1)=0$ с производной $1$ ; $\operatorname{SiLU}(0)=0$ с производной $1/2$ .
Вычитание максимума предотвращает переполнение, которого можно избежать, и не меняет распределение. Каждая операция log-softmax или среднего NLL по индексам целевых классов за один вызов прямого прохода с проверкой входных данных формирует основной результат и значения вероятностей для своего VJP. Сохранённый тензор содержит побитово те же значения f64, которые были сформированы при прямом проходе; речь идёт о совпадении сохранённых чисел с плавающей запятой, а не о точных вещественных вероятностях. VJP использует эти значения, поэтому softmax повторно не вызывается. «Один вызов» не означает одно чтение логита: сначала отдельный обход проверяет конечность значений, затем RowStats разными обходами находит максимум и сумму сдвинутых экспонент, после чего код вывода записывает вероятности. Две ветви — отдельные операции с отдельными вызовами и сохранёнными тензорами; общими остаются только входные логиты.
Публичная точка входа для выбора строк сначала проверяет доступность операнда, ранг таблицы, index_shape и соответствие числа ID этой форме, а затем обходит ID в плоском порядке. Благодаря проверенному плану ядро копирования строк и VJP накопления вкладов используют уже подтверждённые сведения и не обходят исходный запрос повторно. Операция индексированного NLL проверяет ось классов, равенство числа индексов целевых классов числу групп вдоль этой оси, отклоняет нулевое число групп, а затем проверяет индексы целевых классов в порядке следования групп до вычисления функции потерь или сохраняемых вероятностей.
Нет. Каждый вклад уже содержит множитель $1/4$ от средней функции потерь. Обратный проход через выбор строк суммирует три вклада ровно один раз и не применяет ещё одно среднее.

Задайте начальные значения параметров, которые будут обновляться по этим градиентам

К этому моменту реализация курса умеет дифференцировать компактную цепочку от выбранных строк эмбеддингов через проекцию и нелинейную функцию до устойчивой средней функции потерь по токенам. Это пример на уровне отдельных операций, а не архитектура итогового декодера: в последующих блоках SiLU находится внутри сети прямого распространения, а логиты для функции потерь создаёт отдельная проекция в словарь.

В следующих главах из этих операций будут собраны слои эмбеддингов, линейные слои, слои нормализации, SwiGLU и механизм внимания. Правильные градиенты ещё не задают полезные начальные значения параметров. Поэтому в главе 17 появится детерминированная и учитывающая масштаб инициализация, которая нарушает симметрию и не требует нового VJP.