← All chapters

20 · Content revision 3

Let one learned branch gate another

Build a position-wise SwiGLU feed-forward layer, follow its activated gate and linear up branches, and verify exact outputs and gradients.

Predict the two branch products

Chapter 19 supplied reusable bias-free projections. This fixture sends two position rows through three shared weights:

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}

The first row of XWgXW_g is [1,0,1][-1,0,1]. Before looking ahead, predict SiLU([1,0,1])\operatorname{SiLU}([-1,0,1]), multiply it coordinate by coordinate with the first row of XWu=[1,2,3]XW_u=[1,2,3], and then apply W2W_2.

The rounded stages are:

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}

The executable fixture enters the cumulative layer once and retains those exact intermediates for the trace:

Evaluate the fixed two-position SwiGLU example 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)?;

Activate one branch, then multiply

The shared forward formula is:

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

SiLU uses a sigmoid internally but does not return a probability:

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

Thus SiLU(1)<0\operatorname{SiLU}(-1)<0, SiLU(0)=0\operatorname{SiLU}(0)=0, and SiLU(z)z\operatorname{SiLU}(z)\approx z for a large positive zz. Calling the whole activated branch a value between zero and one would be a misconception.

For the reverse path, name the forward intermediates explicitly:

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.

Let G=L/YG=\partial L/\partial Y be the gradient of scalar loss LL, and use dAdA, dSdS, dUdU, dHdH, dXdX, and dWdW for reverse gradients with respect to the corresponding forward quantities. The index pp ranges over every preserved leading position. Reverse mode first crosses the down projection, then splits at the product and crosses the SiLU derivative:

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}

Here SiLU(a)=σ(a)+aσ(a)(1σ(a))\operatorname{SiLU}'(a)=\sigma(a)+a\sigma(a)(1-\sigma(a)). The diagram’s gradient before SiLU is dApdA_p. Forward results and dXpdX_p stay local to position pp. The three weight gradients sum evidence from every position because the weights are shared.

Expand features without mixing positions

  • XX has shape [,din][\ldots,d_{in}]; each leading coordinate is one independent position.
  • WgW_g and WuW_u both have shape [din,dff][d_{in},d_{ff}].
  • W2W_2 has shape [dff,dout][d_{ff},d_{out}].
  • SiLU(z)\operatorname{SiLU}(z) is the elementwise product zσ(z)z\sigma(z).
  • \odot multiplies equal branch coordinates without a matrix product.
  • dind_{in}, dffd_{ff}, and doutd_{out} are input, branch, and output feature widths.
  • \ldots denotes any preserved leading axes.

The worked layer expands 232\to3 and contracts 323\to2. Inputs shaped [2][2], [2,2][2,2], and [1,2,2][1,2,2] produce outputs with those same shapes because this fixture chooses dout=din=2d_{out}=d_{in}=2. The reusable module also permits a different doutd_{out}; the later residual-wrapped decoder chooses equal input and output widths deliberately.

From nonlinear neural language models to SwiGLU

An early feed-forward neural language model used one elementwise tanh hidden transformation over a fixed context. The original Transformer made the feed-forward computation position-wise and wider, but its single activated branch still lacked an input-dependent multiplicative interaction between two learned projections.

Bengio et al., A Neural Probabilistic Language Model: Bengio et al. place an elementwise hyperbolic-tangent hidden layer between learned context word features and next-word scores in a feed-forward neural language model.

Their published score computation includes:

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

The Transformer applies two learned transformations with ReLU separately and identically at every position. Shazeer then tests GLU-family replacements whose two projected branches meet through elementwise multiplication; the SwiGLU variant activates one branch with Swish at beta one before the product.

Vaswani et al., Attention Is All You Need: Vaswani et al. apply the same two-transformation ReLU feed-forward network separately at every sequence position and expand from model width 512 to inner width 2048 in the published configuration.

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 defines bias-free Transformer GLU variants with three weight matrices and writes SwiGLU as a Swish-activated projection multiplied elementwise by a second projection before the output projection.

In Shazeer’s notation, Swish1(z)=zσ(z)\operatorname{Swish}_1(z)=z\sigma(z); this course uses the equivalent name SiLU(z)\operatorname{SiLU}(z) and applies the same position-wise map across every preserved leading position.

A modern decoder can use a bias-free SwiGLU sublayer to expand each token representation, modulate one learned branch with another, and contract to the width needed by the next residual path, while attention remains responsible for mixing positions.

This is the road from nonlinear neural-language-model computation to modern LLM feed-forward blocks, not a programming-language history. The paper reports results for specific tested configurations; those experiments do not by themselves establish why SwiGLU works. The papers establish the architecture, not the exact dimensions, bias policy, parameter names, seed, or error behavior used by this implementation.

The historical Rust sample evaluates tanh and ReLU on the same three inputs; the comparison is about activations, while Rust is only the executable medium:

Compare tanh and ReLU on negative, zero, and positive inputs 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()
}

Compose the cumulative differentiable operations

Construction and forwarding keep projection-stage and branch-dimension failures explicit:

Attribute failures to a projection or composed operation 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 owns gate, up, and down Linear modules. All three disable bias. The forward pass calls the existing projection, SiLU, and multiplication operations; it does not add a fused tensor implementation or another gradient engine:

Compose three bias-free projections around a SiLU gate 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
    }
}

The identity upstream seed exposes both reverse branches and all shared weight gradients. Small parameter-leaf probes expose local intermediate VJPs without changing the layer computation:

Reverse the exact fixture through both branches 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");

Initialization uses one trial stream in gate, up, down order. Equal seeds reproduce values, failed aggregate construction leaves the caller’s stream unchanged, and clones preserve leaf identity:

Initialize all three weights reproducibly and transactionally 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()));

Replacing position zero with an all-zero row proves the forward boundary:

Change one input position and preserve the other output 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;

The learner report covers exact stages, activations, reverse values, shapes, parameters, initialization, identity, independence, empty input, and errors:

Prepare the deterministic Chapter 20 learner report rust/demos/ch20-swiglu-feed-forward/src/main.rs#learner-swiglu-output
    let report = learner_report()?;

Run cargo run --quiet --locked -p ch20-swiglu-feed-forward to inspect the deterministic report. The implementation also covers stable SiLU limits, exact parameter order and count, validation precedence, sequential random draws, all rank variants, and sampled finite differences with step 10610^{-6} and tolerance 3×1063\times10^{-6}.

Trace the gate, merge, and reverse split

The eleven-line trace records two forward rows, two reverse rows, three shared weight gradients, and one position-independence probe:

Emit the exact branch and gradient evidence 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)
}

Follow two projected branches through one position-wise SwiGLU layer

Follow Rust-authored forward values, branch gradients, shared-parameter sums, and a position-independence probe for one two-position fixture.

Input width
22
Branch width
33
Output width
22
Projection policy
Bias-free
Parameter scalars
1818
Input shape
[2,2]\left[2,2\right]
Branch shape
[2,3]\left[2,3\right]
Output shape
[2,2]\left[2,2\right]

Transform each position independently

Transform each position independently
Computation Position 00 Position 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]

The same three weight matrices serve both rows; positions never connect.

SiLU can be negative, zero, or positive; it is not a probability mask.

The diagram rounds to six decimals; the report keeps twelve.

Change one position; observe the other

Rust replaces position zero with an all-zero input and recomputes the layer. Position one is byte-for-byte unchanged.

Changed position
00
Replacement input
[0,0]\left[0,0\right]
Observed position
11
Output before
[0.268941,1.731059]\left[-0.268941,1.731059\right]
Output after
[0.268941,1.731059]\left[-0.268941,1.731059\right]
Result
Unchanged

Split local gradients; accumulate shared weights

Each input receives a local gradient. The three shared weights collect contributions from both positions.

Gradients returned through the two branches at each position: Position 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]
Gradients returned through the two branches at each position: Position 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]
Gradients accumulated into the three shared weights
θ\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}

The figure keeps each position’s two forward branches together, then separates local reverse gradients from the sums accumulated into shared weights. The independence probe changes only position zero, so the unchanged output at position one is direct evidence that SwiGLU transforms positions independently.

Predict before checking the executable evidence

  1. Predict SiLU(1)\operatorname{SiLU}(-1), SiLU(0)\operatorname{SiLU}(0), and SiLU(1)\operatorname{SiLU}(1).
  2. Predict h0h_0 after multiplying the first activated gate by u0=[1,2,3]u_0=[1,2,3].
  3. Predict the three parameter shapes and total scalar count.
  4. Decide which output rows change when only position zero becomes [0,0][0,0].
  5. Decide whether dWgdW_g, dWudW_u, and dW2dW_2 use one position or both.
  6. Predict output shapes for inputs [2][2], [2,2][2,2], and [1,2,2][1,2,2].
  7. Explain why two gated projections cannot generally collapse into one fixed matrix.
  8. Match Bengio, Vaswani, and Shazeer to tanh neural LM, position-wise ReLU FFN, and SwiGLU.
Check the predictions
  1. The values are approximately [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. The shapes are [2,3][2,3], [2,3][2,3], and [3,2][3,2], for 18 scalars and no biases.
  4. Only output row zero changes; output row one is exactly unchanged.
  5. All three shared-weight gradients accumulate contributions from both positions.
  6. With dout=2d_{out}=2, the output shapes remain [2][2], [2,2][2,2], and [1,2,2][1,2,2].
  7. The product makes the effective response of one branch depend on the other branch’s input-dependent values.
  8. Bengio supplies the tanh neural-language-model context, Vaswani the position-wise ReLU block, and Shazeer the later SwiGLU formula.

Hand the nonlinear token transform to batching

The cumulative model gains a named differentiable position-wise SwiGLU module with stable bias-free parameters. Later decoder blocks will choose equal input and output widths so this sublayer can sit inside a residual path; Chapter 21 first defines how losses and gradients from multiple token examples combine.

SwiGLU changes features within one token position. It does not move information between tokens; later causal attention owns that job. Chapter 21 now groups causal examples and specifies how token losses and their gradients are averaged.