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:
The first row of is . Before looking ahead, predict , multiply it coordinate by coordinate with the first row of , and then apply .
The rounded stages are:
The executable fixture enters the cumulative layer once and retains those exact intermediates for the trace:
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:
SiLU uses a sigmoid internally but does not return a probability:
Thus , , and for a large positive . Calling the whole activated branch a value between zero and one would be a misconception.
For the reverse path, name the forward intermediates explicitly:
Let be the gradient of scalar loss , and use , , , , , and for reverse gradients with respect to the corresponding forward quantities. The index ranges over every preserved leading position. Reverse mode first crosses the down projection, then splits at the product and crosses the SiLU derivative:
Here . The diagram’s gradient before SiLU is . Forward results and stay local to position . The three weight gradients sum evidence from every position because the weights are shared.
Expand features without mixing positions
- has shape ; each leading coordinate is one independent position.
- and both have shape .
- has shape .
- is the elementwise product .
- multiplies equal branch coordinates without a matrix product.
- , , and are input, branch, and output feature widths.
- denotes any preserved leading axes.
The worked layer expands and contracts . Inputs shaped , , and produce outputs with those same shapes because this fixture chooses . The reusable module also permits a different ; 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:
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.
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, ; this course uses the equivalent name 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:
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:
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:
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:
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:
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:
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:
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 and tolerance
.
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:
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
- Branch width
- Output width
- Projection policy
- Bias-free
- Parameter scalars
- Input shape
- Branch shape
- Output shape
Transform each position independently
| Computation | Position | Position |
|---|---|---|
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
- Replacement input
- Observed position
- Output before
- Output after
- Result
- Unchanged
Split local gradients; accumulate shared weights
Each input receives a local gradient. The three shared weights collect contributions from both positions.
ffn.gate.weight | |||
|---|---|---|---|
ffn.up.weight | |||
ffn.down.weight |
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
- Predict , , and .
- Predict after multiplying the first activated gate by .
- Predict the three parameter shapes and total scalar count.
- Decide which output rows change when only position zero becomes .
- Decide whether , , and use one position or both.
- Predict output shapes for inputs , , and .
- Explain why two gated projections cannot generally collapse into one fixed matrix.
- Match Bengio, Vaswani, and Shazeer to tanh neural LM, position-wise ReLU FFN, and SwiGLU.
Check the predictions
- The values are approximately .
- .
- The shapes are , , and , for 18 scalars and no biases.
- Only output row zero changes; output row one is exactly unchanged.
- All three shared-weight gradients accumulate contributions from both positions.
- With , the output shapes remain , , and .
- The product makes the effective response of one branch depend on the other branch’s input-dependent values.
- 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.