← Все главы

20 · Версия материала 3

Управляйте ветвью расширения с помощью обучаемого вентиля

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

Предскажите произведения двух ветвей

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

X=[1001],Wg=[101011],Wu=[123321],W2=[100111].\begin{aligned} X &= \begin{bmatrix}1&0\\0&1\end{bmatrix}, \\ W_g &= \begin{bmatrix}-1&0&1\\0&1&-1\end{bmatrix}, \\ W_u &= \begin{bmatrix}1&2&3\\3&2&1\end{bmatrix}, \\ W_2 &= \begin{bmatrix}1&0\\0&1\\1&-1\end{bmatrix}. \end{aligned}

Первая строка XWgXW_g равна [1,0,1][-1,0,1]. Прежде чем читать дальше, предскажите SiLU([1,0,1])\operatorname{SiLU}([-1,0,1]), поэлементно умножьте результат на первую строку XWu=[1,2,3]XW_u=[1,2,3], а затем примените W2W_2.

После округления этапы выглядят так:

XWg=[101011],SiLU(XWg)[0.26894100.73105900.7310590.268941],XWu=[123321],SiLU(XWg)(XWu)[0.26894102.19317601.4621170.268941],Y[1.9242342.1931760.2689411.731059].\begin{aligned} XW_g &= \begin{bmatrix}-1&0&1\\0&1&-1\end{bmatrix}, \\ \operatorname{SiLU}(XW_g) &\approx \begin{bmatrix}-0.268941&0&0.731059\\0&0.731059&-0.268941\end{bmatrix}, \\ XW_u &= \begin{bmatrix}1&2&3\\3&2&1\end{bmatrix}, \\ \operatorname{SiLU}(XW_g)\odot(XW_u) &\approx \begin{bmatrix}-0.268941&0&2.193176\\0&1.462117&-0.268941\end{bmatrix}, \\ Y &\approx \begin{bmatrix}1.924234&-2.193176\\-0.268941&1.731059\end{bmatrix}. \end{aligned}

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

Вычислить фиксированный пример SwiGLU для двух позиций rust/demos/ch20-swiglu-feed-forward/src/lib.rs#known-swiglu-forward
    let layer = known_swiglu();
    let input = TensorValue::parameter(tensor(&INPUT_SHAPE, &INPUT_VALUES))?;
    let pass = layer.forward_with_intermediates(&input)?;

Активируйте одну ветвь, затем перемножьте

Общая формула прямого прохода:

FFN(X)=(SiLU(XWg)(XWu))W2\operatorname{FFN}(X)=\left(\operatorname{SiLU}(XW_g)\odot(XW_u)\right)W_2

Внутри SiLU используется сигмоида, но результат активации не является вероятностью:

σ(z)=11+ez,SiLU(z)=zσ(z).\sigma(z)=\frac{1}{1+e^{-z}}, \qquad \operatorname{SiLU}(z)=z\sigma(z).

Поэтому SiLU(1)<0\operatorname{SiLU}(-1)<0, SiLU(0)=0\operatorname{SiLU}(0)=0, а при большом положительном zz выполняется SiLU(z)z\operatorname{SiLU}(z)\approx z. Считать, что вся активированная ветвь принимает значения только от нуля до единицы, было бы ошибкой.

Для обратного прохода явно обозначим промежуточные результаты прямого прохода:

A=XWg,S=SiLU(A),U=XWu,H=SU,Y=HW2.A=XW_g,\qquad S=\operatorname{SiLU}(A),\qquad U=XW_u,\qquad H=S\odot U,\qquad Y=HW_2.

Пусть G=L/YG=\partial L/\partial Y — градиент скалярной функции потерь LL, а dAdA, dSdS, dUdU, dHdH, dXdX и dWdW обозначают градиенты по соответствующим величинам прямого прохода. Индекс pp перебирает все сохранённые ведущие позиции. При обратном проходе сначала вычисляются градиенты через проекцию сжатия, затем градиент разветвляется в узле произведения и проходит через производную SiLU:

dH=GW2,dS=dHU,dU=dHS,dA=dSSiLU(A),dXp=dApWg+dUpWu,dWg=pXpdAp,dWu=pXpdUp,dW2=pHpGp.\begin{aligned} dH &= GW_2^\top, \\ dS &= dH\odot U, \\ dU &= dH\odot S, \\ dA &= dS\odot\operatorname{SiLU}'(A), \\ dX_p &= dA_pW_g^\top+dU_pW_u^\top, \\ dW_g &= \sum_p X_p^\top dA_p, \\ dW_u &= \sum_p X_p^\top dU_p, \\ dW_2 &= \sum_p H_p^\top G_p. \end{aligned}

Здесь SiLU(a)=σ(a)+aσ(a)(1σ(a))\operatorname{SiLU}'(a)=\sigma(a)+a\sigma(a)(1-\sigma(a)). В схеме градиент до SiLU обозначен dApdA_p. Результаты прямого прохода и dXpdX_p остаются локальными для позиции pp. В градиентах трёх матриц весов суммируются вклады всех позиций, поскольку эти веса общие.

Расширяйте признаки, не смешивая позиции

  • XX имеет форму [,din][\ldots,d_{in}]; каждая ведущая координата задаёт отдельную позицию.
  • WgW_g и WuW_u имеют форму [din,dff][d_{in},d_{ff}].
  • W2W_2 имеет форму [dff,dout][d_{ff},d_{out}].
  • SiLU(z)\operatorname{SiLU}(z) — поэлементное произведение zσ(z)z\sigma(z).
  • \odot перемножает соответствующие координаты ветвей без матричного умножения.
  • dind_{in}, dffd_{ff} и doutd_{out} — ширины входных признаков, ветвей и выходных признаков.
  • \ldots обозначает любые сохраняемые ведущие оси.

Слой из примера расширяет признаки как 232\to3 и сжимает как 323\to2. Входы форм [2][2], [2,2][2,2] и [1,2,2][1,2,2] дают выходы тех же форм, потому что в этом примере выбрано dout=din=2d_{out}=d_{in}=2. Этот модуль допускает и другие значения doutd_{out}; позже в декодере с остаточными связями ширины входа и выхода будут намеренно одинаковыми.

От нелинейных нейросетевых моделей языка к SwiGLU

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

Бенжио и соавторы, A Neural Probabilistic Language Model: Бенжио и соавторы помещают скрытый слой с поэлементным гиперболическим тангенсом между обучаемыми признаками слов контекста и оценками следующего слова в нейросетевой модели языка с прямым распространением.

Опубликованное ими вычисление оценок включает:

y=b+Wx+Utanh(d+Hx).y=b+Wx+U\tanh(d+Hx).

Transformer применяет два обучаемых преобразования с ReLU отдельно и одинаково в каждой позиции. Затем Shazeer исследует замены из семейства GLU, в которых две спроецированные ветви соединяются поэлементным умножением; в варианте SwiGLU одна ветвь перед произведением активируется функцией Swish при β=1\beta=1.

Васвани и соавторы, Attention Is All You Need: Васвани и соавторы применяют одну и ту же сеть прямого распространения с двумя преобразованиями и ReLU отдельно к каждой позиции последовательности; в опубликованной конфигурации ширина модели 512512 расширяется до внутренней ширины 20482048.

FFN(x)=max(0,xW1+b1)W2+b2.\operatorname{FFN}(x)=\max(0,xW_1+b_1)W_2+b_2.

Shazeer, GLU Variants Improve Transformer: Shazeer определяет варианты GLU для Transformer без смещений с тремя матрицами весов и записывает SwiGLU как проекцию с активацией Swish, поэлементно умноженную на вторую проекцию перед выходной проекцией.

В обозначениях Shazeer Swish1(z)=zσ(z)\operatorname{Swish}_1(z)=z\sigma(z); в этом курсе используется равнозначное название SiLU(z)\operatorname{SiLU}(z), а одно и то же преобразование независимо применяется ко всем сохранённым ведущим позициям.

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

Так проходит путь от нелинейных вычислений в нейросетевых моделях языка к современным блокам прямого распространения LLM, а не история языков программирования. В работе Shazeer приведены результаты для конкретных исследованных конфигураций, но сами по себе эти эксперименты не объясняют, почему SwiGLU работает. Источники определяют архитектуру, но не размеры, вариант со смещениями или без них, имена параметров, начальное состояние генератора и обработку ошибок в этой реализации.

Исторический пример на Rust вычисляет tanh и ReLU для одних и тех же трёх входов. Сравниваются функции активации, а код на Rust лишь позволяет исполнить этот пример:

Сравнить tanh и ReLU на отрицательном, нулевом и положительном входах rust/demos/ch20-swiglu-feed-forward/src/lib.rs#historical-activation-contrast
/// Evaluates the tanh hidden activation used by an early neural language model.
pub fn tanh_hidden(values: &[f64]) -> Vec<f64> {
    values.iter().map(|value| value.tanh()).collect()
}

/// Evaluates the ReLU used by the original Transformer's position-wise FFN.
pub fn relu_hidden(values: &[f64]) -> Vec<f64> {
    values.iter().map(|value| value.max(0.0)).collect()
}

Скомпонуйте уже реализованные дифференцируемые операции

Ошибки создания и прямого прохода указывают конкретный этап проекции или несовместимые размеры ветвей:

Связать каждую ошибку с проекцией или составной операцией rust/crates/llm-from-scratch/src/nn/swiglu.rs#swiglu-errors
/// A rejected parameter set, input, or delegated differentiable operation.
#[derive(Clone, Debug, PartialEq)]
pub enum SwiGluError {
    Projection {
        projection: SwiGluProjection,
        source: LinearError,
    },
    Autodiff {
        operation: SwiGluOperation,
        source: TensorAutodiffError,
    },
    BranchInputWidthMismatch {
        gate: usize,
        up: usize,
    },
    BranchHiddenWidthMismatch {
        gate: usize,
        up: usize,
    },
    DownInputWidthMismatch {
        hidden: usize,
        down: usize,
    },
    Initialization(InitializationError),
}

impl fmt::Display for SwiGluError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Projection { projection, source } => {
                write!(formatter, "SwiGLU {projection} projection: {source}")
            }
            Self::Autodiff { operation, source } => {
                write!(formatter, "SwiGLU {operation}: {source}")
            }
            Self::BranchInputWidthMismatch { gate, up } => write!(
                formatter,
                "SwiGLU gate and up input widths must match, got {gate} and {up}"
            ),
            Self::BranchHiddenWidthMismatch { gate, up } => write!(
                formatter,
                "SwiGLU gate and up hidden widths must match, got {gate} and {up}"
            ),
            Self::DownInputWidthMismatch { hidden, down } => write!(
                formatter,
                "SwiGLU down input width must equal hidden width {hidden}, got {down}"
            ),
            Self::Initialization(source) => source.fmt(formatter),
        }
    }
}

impl Error for SwiGluError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Projection { source, .. } => Some(source),
            Self::Autodiff { source, .. } => Some(source),
            Self::Initialization(source) => Some(source),
            _ => None,
        }
    }
}

SwiGlu содержит вентильный модуль Linear, модуль расширения и модуль сжатия. Во всех трёх смещение отключено. Прямой проход вызывает уже реализованные проекцию, SiLU и умножение; отдельная слитная тензорная операция или ещё один механизм вычисления градиентов здесь не появляются:

Скомпоновать три проекции без смещения вокруг вентиля SiLU rust/crates/llm-from-scratch/src/nn/swiglu.rs#swiglu-layer
/// The exact tensors produced by one composed SwiGLU forward pass.
#[derive(Clone, Debug)]
pub struct SwiGluForward {
    gate_linear: TensorValue,
    gate_silu: TensorValue,
    up: TensorValue,
    product: TensorValue,
    output: TensorValue,
}

impl SwiGluForward {
    pub fn gate_linear(&self) -> &TensorValue {
        &self.gate_linear
    }

    pub fn gate_silu(&self) -> &TensorValue {
        &self.gate_silu
    }

    pub fn up(&self) -> &TensorValue {
        &self.up
    }

    pub fn product(&self) -> &TensorValue {
        &self.product
    }

    pub fn output(&self) -> &TensorValue {
        &self.output
    }

    pub fn into_output(self) -> TensorValue {
        self.output
    }
}

/// Three bias-free projections with a SiLU-activated multiplicative gate.
#[derive(Clone, Debug)]
pub struct SwiGlu {
    gate: Linear,
    up: Linear,
    down: Linear,
    parameters: NamedParameters,
    input_width: usize,
    hidden_width: usize,
    output_width: usize,
}

impl SwiGlu {
    /// Initializes all three projections transactionally from one deterministic stream.
    pub fn new(
        parameter_prefix: impl Into<String>,
        input_width: usize,
        hidden_width: usize,
        output_width: usize,
        rng: &mut SplitMix64,
    ) -> Result<Self, SwiGluError> {
        let parameter_prefix = parameter_prefix.into();
        let mut trial = rng.clone();
        let gate = Linear::new(
            format!("{parameter_prefix}.gate"),
            input_width,
            hidden_width,
            false,
            &mut trial,
        )
        .map_err(projection_error(SwiGluProjection::Gate))?;
        let up = Linear::new(
            format!("{parameter_prefix}.up"),
            input_width,
            hidden_width,
            false,
            &mut trial,
        )
        .map_err(projection_error(SwiGluProjection::Up))?;
        let down = Linear::new(
            format!("{parameter_prefix}.down"),
            hidden_width,
            output_width,
            false,
            &mut trial,
        )
        .map_err(projection_error(SwiGluProjection::Down))?;
        let layer = Self::from_linears(gate, up, down)?;
        *rng = trial;
        Ok(layer)
    }

    /// Gives SwiGLU semantics to three existing bias-free weight matrices.
    pub fn from_parameters(
        gate_weight: NamedParameter,
        up_weight: NamedParameter,
        down_weight: NamedParameter,
    ) -> Result<Self, SwiGluError> {
        let gate = Linear::from_parameters(gate_weight, None)
            .map_err(projection_error(SwiGluProjection::Gate))?;
        let up = Linear::from_parameters(up_weight, None)
            .map_err(projection_error(SwiGluProjection::Up))?;
        let down = Linear::from_parameters(down_weight, None)
            .map_err(projection_error(SwiGluProjection::Down))?;
        Self::from_linears(gate, up, down)
    }

    fn from_linears(gate: Linear, up: Linear, down: Linear) -> Result<Self, SwiGluError> {
        if gate.input_width() != up.input_width() {
            return Err(SwiGluError::BranchInputWidthMismatch {
                gate: gate.input_width(),
                up: up.input_width(),
            });
        }
        if gate.output_width() != up.output_width() {
            return Err(SwiGluError::BranchHiddenWidthMismatch {
                gate: gate.output_width(),
                up: up.output_width(),
            });
        }
        if down.input_width() != gate.output_width() {
            return Err(SwiGluError::DownInputWidthMismatch {
                hidden: gate.output_width(),
                down: down.input_width(),
            });
        }
        let parameters = NamedParameters::try_new(
            gate.parameters()
                .iter()
                .chain(up.parameters())
                .chain(down.parameters())
                .cloned()
                .collect(),
        )
        .map_err(SwiGluError::Initialization)?;
        let input_width = gate.input_width();
        let hidden_width = gate.output_width();
        let output_width = down.output_width();
        Ok(Self {
            gate,
            up,
            down,
            parameters,
            input_width,
            hidden_width,
            output_width,
        })
    }

    /// Applies the same gated feature transformation at every leading position.
    pub fn forward(&self, input: &TensorValue) -> Result<TensorValue, SwiGluError> {
        Ok(self.forward_with_intermediates(input)?.into_output())
    }

    /// Returns each branch tensor for inspection without changing the computation.
    pub fn forward_with_intermediates(
        &self,
        input: &TensorValue,
    ) -> Result<SwiGluForward, SwiGluError> {
        let gate_linear = self
            .gate
            .forward(input)
            .map_err(projection_error(SwiGluProjection::Gate))?;
        let gate_silu = gate_linear
            .silu()
            .map_err(autodiff_error(SwiGluOperation::SiluGate))?;
        let up = self
            .up
            .forward(input)
            .map_err(projection_error(SwiGluProjection::Up))?;
        let product = gate_silu
            .mul(&up)
            .map_err(autodiff_error(SwiGluOperation::ElementwiseGate))?;
        let output = self
            .down
            .forward(&product)
            .map_err(projection_error(SwiGluProjection::Down))?;
        Ok(SwiGluForward {
            gate_linear,
            gate_silu,
            up,
            product,
            output,
        })
    }

    pub fn gate(&self) -> &Linear {
        &self.gate
    }

    pub fn up(&self) -> &Linear {
        &self.up
    }

    pub fn down(&self) -> &Linear {
        &self.down
    }

    pub fn parameters(&self) -> &[NamedParameter] {
        self.parameters.as_slice()
    }

    pub const fn input_width(&self) -> usize {
        self.input_width
    }

    pub const fn hidden_width(&self) -> usize {
        self.hidden_width
    }

    pub const fn output_width(&self) -> usize {
        self.output_width
    }

    pub const fn parameter_count(&self) -> usize {
        2 * self.input_width * self.hidden_width + self.hidden_width * self.output_width
    }
}

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

Провести точный пример назад через обе ветви rust/demos/ch20-swiglu-feed-forward/src/lib.rs#swiglu-gradients
    let upstream = tensor(&UPSTREAM_SHAPE, &UPSTREAM_VALUES);
    pass.output()
        .backward_with_seed(&upstream.view(), GraphRetention::Retain)?;
    let input_gradient = input
        .gradient_snapshot()
        .expect("trainable input stores its reverse gradient");
    let gate_weight_gradient = layer
        .gate()
        .weight()
        .tensor()
        .gradient_snapshot()
        .expect("gate dW");
    let up_weight_gradient = layer
        .up()
        .weight()
        .tensor()
        .gradient_snapshot()
        .expect("up dW");
    let down_weight_gradient = layer
        .down()
        .weight()
        .tensor()
        .gradient_snapshot()
        .expect("down dW");

    // Persistent gradients belong to parameter leaves. These two tiny probes
    // promote recorded intermediates to leaves so their local VJPs are visible.
    let product_probe_layer = known_swiglu();
    let product_probe = TensorValue::parameter(pass.product().value_snapshot())?;
    product_probe_layer
        .down()
        .forward(&product_probe)?
        .backward_with_seed(&upstream.view(), GraphRetention::Release)?;
    let product_gradient = product_probe
        .gradient_snapshot()
        .expect("product probe stores its reverse gradient");

    let branch_probe_layer = known_swiglu();
    let gate_linear_probe = TensorValue::parameter(pass.gate_linear().value_snapshot())?;
    let up_probe = TensorValue::parameter(pass.up().value_snapshot())?;
    let branch_product = gate_linear_probe.silu()?.mul(&up_probe)?;
    branch_probe_layer
        .down()
        .forward(&branch_product)?
        .backward_with_seed(&upstream.view(), GraphRetention::Release)?;
    let gate_linear_gradient = gate_linear_probe
        .gradient_snapshot()
        .expect("gate probe stores its reverse gradient");
    let up_gradient = up_probe
        .gradient_snapshot()
        .expect("up probe stores its reverse gradient");

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

Воспроизводимо инициализировать все три матрицы весов без частичного изменения состояния rust/demos/ch20-swiglu-feed-forward/src/lib.rs#initialized-swiglu
    let mut first_rng = SplitMix64::from_seed(20);
    let mut second_rng = SplitMix64::from_seed(20);
    let initialized = SwiGlu::new("ffn", 2, 3, 2, &mut first_rng)?;
    let reproduced = SwiGlu::new("ffn", 2, 3, 2, &mut second_rng)?;
    let initialized_reproducible = initialized
        .parameters()
        .iter()
        .zip(reproduced.parameters())
        .all(|(left, right)| *left.tensor().value() == *right.tensor().value());
    let cloned = initialized.clone();
    let clone_shares_parameters = initialized
        .parameters()
        .iter()
        .zip(cloned.parameters())
        .all(|(left, right)| left.tensor().is_same_node(right.tensor()));

Если заменить вход позиции 0 полностью нулевой строкой, становится видно, что прямой проход не смешивает позиции:

Изменить одну позицию входа и сохранить выход другой rust/demos/ch20-swiglu-feed-forward/src/lib.rs#position-independence
    let perturbed = layer
        .forward(&TensorValue::constant(tensor(
            &INPUT_SHAPE,
            &PERTURBED_INPUT_VALUES,
        ))?)?
        .value_snapshot();
    let independent_before = pass.output().value().as_slice()[2..4].to_vec();
    let independent_after = perturbed.as_slice()[2..4].to_vec();
    let position_independent = independent_before == independent_after;

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

Подготовить детерминированный учебный отчёт главы 20 rust/demos/ch20-swiglu-feed-forward/src/main.rs#learner-swiglu-output
    let report = learner_report()?;

Команда cargo run --quiet --locked -p ch20-swiglu-feed-forward выводит детерминированный отчёт. Реализация также проверяет устойчивые предельные значения SiLU, точный порядок и число параметров, очерёдность проверок, последовательные случайные выборки, все варианты ранга и выборочную проверку конечными разностями с шагом 10610^{-6} и допуском 3×1063\times10^{-6}.

Проследите вентиль, слияние и разветвление обратного прохода

Трасса из одиннадцати строк содержит две строки прямого прохода, две строки обратного прохода, три градиента общих матриц весов и одну проверку независимости позиций:

Вывести точные значения ветвей и градиентов rust/demos/ch20-swiglu-feed-forward/src/diagram_trace.rs#swiglu-feed-forward-trace
/// Renders the exact Rust-owned evidence consumed by the static chapter diagram.
pub fn render_trace() -> Result<String, Box<dyn Error>> {
    let report = learner_report()?;
    let (position_count, model_width) = matrix_dimensions(report.input.shape(), "input")?;
    let (branch_positions, hidden_width) =
        matrix_dimensions(report.product.shape(), "branch product")?;
    let (output_positions, output_width) = matrix_dimensions(report.output.shape(), "output")?;
    let (upstream_positions, upstream_width) =
        matrix_dimensions(report.upstream.shape(), "upstream")?;
    if branch_positions != position_count
        || output_positions != position_count
        || upstream_positions != position_count
        || upstream_width != output_width
        || report.gate_linear.shape() != report.product.shape()
        || report.gate_silu.shape() != report.product.shape()
        || report.up.shape() != report.product.shape()
        || report.product_gradient.shape() != report.product.shape()
        || report.gate_linear_gradient.shape() != report.product.shape()
        || report.up_gradient.shape() != report.product.shape()
        || report.input_gradient.shape() != report.input.shape()
    {
        return Err("learner report shapes disagree with the position-wise trace".into());
    }
    let mut trace = String::new();
    writeln!(trace, "TRACE swiglu-feed-forward-v1 BEGIN")?;
    writeln!(
        trace,
        "FIXTURE name=known-position-wise-swiglu model-width={} hidden-width={} output-width={} bias=false parameter-count={} input-shape={} branch-shape={} output-shape={} upstream-shape={}",
        model_width,
        hidden_width,
        output_width,
        report.parameter_count,
        shape(report.input.shape()),
        shape(report.product.shape()),
        shape(report.output.shape()),
        shape(report.upstream.shape())
    )?;
    for position in 0..position_count {
        let input_start = position * model_width;
        let hidden_start = position * hidden_width;
        let output_start = position * output_width;
        writeln!(
            trace,
            "POSITION-FORWARD position={} input={} gate-pre={} gate-silu={} up={} gated={} output={}",
            position,
            fixed_list(&report.input.as_slice()[input_start..input_start + model_width]),
            fixed_list(&report.gate_linear.as_slice()[hidden_start..hidden_start + hidden_width]),
            fixed_list(&report.gate_silu.as_slice()[hidden_start..hidden_start + hidden_width]),
            fixed_list(&report.up.as_slice()[hidden_start..hidden_start + hidden_width]),
            fixed_list(&report.product.as_slice()[hidden_start..hidden_start + hidden_width]),
            fixed_list(&report.output.as_slice()[output_start..output_start + output_width])
        )?;
    }
    for position in 0..position_count {
        let input_start = position * model_width;
        let hidden_start = position * hidden_width;
        let upstream_start = position * upstream_width;
        writeln!(
            trace,
            "POSITION-BACKWARD position={} upstream={} gated-gradient={} gate-gradient={} up-gradient={} input-gradient={}",
            position,
            fixed_list(
                &report.upstream.as_slice()[upstream_start..upstream_start + upstream_width]
            ),
            fixed_list(
                &report.product_gradient.as_slice()[hidden_start..hidden_start + hidden_width]
            ),
            fixed_list(
                &report.gate_linear_gradient.as_slice()[hidden_start..hidden_start + hidden_width]
            ),
            fixed_list(&report.up_gradient.as_slice()[hidden_start..hidden_start + hidden_width]),
            fixed_list(&report.input_gradient.as_slice()[input_start..input_start + model_width])
        )?;
    }
    for (parameter, gradient) in report.parameter_names.iter().zip([
        &report.gate_weight_gradient,
        &report.up_weight_gradient,
        &report.down_weight_gradient,
    ]) {
        writeln!(
            trace,
            "PARAMETER-GRADIENT name={} shape={} values={}",
            parameter,
            shape(gradient.shape()),
            fixed_list(gradient.as_slice())
        )?;
    }
    let replacement_input = PERTURBED_INPUT_VALUES
        .get(..model_width)
        .ok_or("perturbed input is shorter than the model width")?;
    writeln!(
        trace,
        "INDEPENDENCE changed-position=0 replacement-input={} observed-position=1 before={} after={} unchanged={}",
        fixed_list(replacement_input),
        fixed_list(&report.independent_before),
        fixed_list(&report.independent_after),
        report.position_independent
    )?;
    writeln!(trace, "TRACE swiglu-feed-forward-v1 END")?;
    Ok(trace)
}

Проследите две ветви проекций в слое SwiGLU, применяемом к каждой позиции

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

Ширина входа
22
Ширина ветвей
33
Ширина выхода
22
Вариант проекций
Без смещений
Скалярных параметров
1818
Форма входа
[2,2]\left[2,2\right]
Форма ветвей
[2,3]\left[2,3\right]
Форма выхода
[2,2]\left[2,2\right]

Обработайте каждую позицию отдельно

Обработайте каждую позицию отдельно
Вычисление Позиция 00 Позиция 11
XX X0=[1,0]X_{0}=\left[1,0\right] X1=[0,1]X_{1}=\left[0,1\right]
g=XWgg=XW_g g0=X0Wgg_{0}=X_{0}W_g g1=X1Wgg_{1}=X_{1}W_g
gg g0=[1,0,1]g_{0}=\left[-1,0,1\right] g1=[0,1,1]g_{1}=\left[0,1,-1\right]
s=SiLU(g)s=\operatorname{SiLU}(g) s0=[0.268941,0,0.731059]s_{0}=\left[-0.268941,0,0.731059\right] s1=[0,0.731059,0.268941]s_{1}=\left[0,0.731059,-0.268941\right]
u=XWuu=XW_u u0=X0Wuu_{0}=X_{0}W_u u1=X1Wuu_{1}=X_{1}W_u
uu u0=[1,2,3]u_{0}=\left[1,2,3\right] u1=[3,2,1]u_{1}=\left[3,2,1\right]
h=suh=s\odot u h0=[0.268941,0,2.193176]h_{0}=\left[-0.268941,0,2.193176\right] h1=[0,1.462117,0.268941]h_{1}=\left[0,1.462117,-0.268941\right]
Y=hW2Y=hW_2 Y0=[1.924234,2.193176]Y_{0}=\left[1.924234,-2.193176\right] Y1=[0.268941,1.731059]Y_{1}=\left[-0.268941,1.731059\right]

Обе строки используют три общие матрицы весов; позиции не связаны.

SiLU даёт число любого знака или ноль; это не вероятностная маска.

Схема округляет до шести знаков; отчёт сохраняет двенадцать.

Измените одну позицию; проверьте другую

Rust заменяет вход позиции 0 нулями и пересчитывает слой; выход позиции 1 побайтно не меняется.

Меняемая позиция
00
Новый вход
[0,0]\left[0,0\right]
Наблюдаемая позиция
11
Выход до замены
[0.268941,1.731059]\left[-0.268941,1.731059\right]
Выход после замены
[0.268941,1.731059]\left[-0.268941,1.731059\right]
Результат
Без изменений

Разделите локальные градиенты и сложите вклады в общие веса

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

Градиенты, возвращаемые через две ветви в каждой позиции: Позиция 00
G0G_{0} [1,0]\left[1,0\right]
dH0dH_{0} [1,0,1]\left[1,0,1\right]
dA0dA_{0} [0.072329,0,2.783012]\left[0.072329,0,2.783012\right]
dU0dU_{0} [0.268941,0,0.731059]\left[-0.268941,0,0.731059\right]
dX0dX_{0} [4.634916,2.858777]\left[4.634916,-2.858777\right]
Градиенты, возвращаемые через две ветви в каждой позиции: Позиция 11
G1G_{1} [0,1]\left[0,1\right]
dH1dH_{1} [0,1,1]\left[0,1,-1\right]
dA1dA_{1} [0,1.855341,0.072329]\left[0,1.855341,-0.072329\right]
dU1dU_{1} [0,0.731059,0.268941]\left[0,0.731059,0.268941\right]
dX1dX_{1} [2.196612,3.658729]\left[2.196612,3.658729\right]
Градиенты, накопленные в трёх общих матрицах весов
θ\theta p\sum_p dim(θ)\dim(\theta) dθd\theta
WgW_gffn.gate.weight pXpdAp\sum_p X_p^\top dA_p [2,3]\left[2,3\right] [0.07232902.78301201.8553410.072329]\begin{bmatrix}0.072329&0&2.783012\\0&1.855341&-0.072329\end{bmatrix}
WuW_uffn.up.weight pXpdUp\sum_p X_p^\top dU_p [2,3]\left[2,3\right] [0.26894100.73105900.7310590.268941]\begin{bmatrix}-0.268941&0&0.731059\\0&0.731059&0.268941\end{bmatrix}
W2W_2ffn.down.weight pHpGp\sum_p H_p^\top G_p [3,2]\left[3,2\right] [0.268941001.4621172.1931760.268941]\begin{bmatrix}-0.268941&0\\0&1.462117\\2.193176&-0.268941\end{bmatrix}

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

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

  1. Предскажите SiLU(1)\operatorname{SiLU}(-1), SiLU(0)\operatorname{SiLU}(0) и SiLU(1)\operatorname{SiLU}(1).
  2. Предскажите h0h_0 после умножения активированной вентильной ветви для позиции 0 на u0=[1,2,3]u_0=[1,2,3].
  3. Предскажите формы трёх матриц параметров и общее число скаляров.
  4. Определите, какие строки выхода изменятся, если только позиция 0 станет равна [0,0][0,0].
  5. Определите, используют ли dWgdW_g, dWudW_u и dW2dW_2 одну позицию или обе.
  6. Предскажите формы выходов для входов [2][2], [2,2][2,2] и [1,2,2][1,2,2].
  7. Объясните, почему такую вентильную обработку обычно нельзя свести к одной постоянной матрице.
  8. Сопоставьте Бенжио, Васвани и Shazeer соответственно с нейросетевой моделью языка с tanh, блоком прямого распространения с ReLU, который одинаково применяется к каждой позиции, и SwiGLU.
Сверить ответы
  1. Приближённые значения равны [0.268941,0,0.731059][-0.268941,0,0.731059].
  2. h0[0.268941,0,2.193176]h_0\approx[-0.268941,0,2.193176].
  3. Формы равны [2,3][2,3], [2,3][2,3] и [3,2][3,2]: всего 18 скаляров и ни одного смещения.
  4. Изменится только строка выхода 0; строка выхода 1 останется в точности прежней.
  5. В градиентах всех трёх общих матриц весов складываются вклады обеих позиций.
  6. При dout=2d_{out}=2 формы выходов остаются равны [2][2], [2,2][2,2] и [1,2,2][1,2,2].
  7. Из-за произведения эффективный отклик одной ветви зависит от значений другой ветви, которые, в свою очередь, зависят от входа.
  8. Бенжио задаёт контекст нейросетевой модели языка с tanh, Васвани — одинаково применяемый к каждой позиции блок с ReLU, а Shazeer — более позднюю формулу SwiGLU.

Перейдите от нелинейного преобразования токенов к пакетной обработке

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

SwiGLU изменяет признаки внутри одной позиции токена, но не переносит информацию между токенами: это позже сделает каузальное внимание. Теперь глава 21 объединит каузальные примеры и точно задаст усреднение функций потерь токенов и их градиентов.