Thread (23 posts)

Great question! You're asking about the order of operations in the FFN variant. Let me clarify the architecture:

FFN GeGLU Structure

From the paper, the FFN variant is: $\text{FFN}_\text{GEGLU}(x, W, V, W_2) = (\text{GELU}(xW) \otimes xV)W_2$

The order of operations is:

1. Compute $\text{GELU}(xW)$ → produces hidden vector
2. Compute $xV$ → produces hidden vector
3. Element-wise multiply: $\text{GELU}(xW) \otimes xV$ → produces hidden vector
4. Finally multiply by $W_2$: $[\text{GELU}(xW) \otimes xV]W_2$ → produces output

Why This Order Matters

Architectural Reason

The multiplication by $W_2$ happens after the gating because:

• The gating mechanism operates in the hidden space (dimension $d_{ff}$)
• $W_2$ projects from hidden space back to model dimension ($d_{model}$)
• This preserves the standard Transformer FFN structure: expand → activate → project

If We Did $(xV)W_2$ First

If we computed $\text{GELU}(xW) \otimes (xVW_2)$, we would:

• Project $xV$ down to model dimension first
• Then apply gating in the model dimension instead of hidden dimension
• This would fundamentally change the architecture and reduce expressiveness

Dimensional Analysis

• $x$: shape $[..., d_{model}]$
• $W, V$: shape $[d_{model}, d_{ff}]$
• $W_2$: shape $[d_{ff}, d_{model}]$
• $\text{GELU}(xW) \otimes xV$: shape $[..., d_{ff}]$ (gating in hidden space)
• Final output: shape $[..., d_{model}]$

The current order maintains the "bottleneck" structure where computation happens in the expanded hidden dimension before projecting back down.

Excellent question! This touches on a fundamental notational convention difference in deep learning. Let me explain the full context.

Notational Conventions

Row Vector Convention (xW)

The paper uses the row vector convention where:

• Input $x$ is a row vector: $x \in \mathbb{R}^{1 \times d_{in}}$
• Weight matrix $W$: $W \in \mathbb{R}^{d_{in} \times d_{ff}}$
• Output: $xW \in \mathbb{R}^{1 \times d_{ff}}$

Column Vector Convention (Wx)

Traditional linear algebra often uses:

• Input $x$ is a column vector: $x \in \mathbb{R}^{d_{in} \times 1}$
• Weight matrix $W$: $W \in \mathbb{R}^{d_{ff} \times d_{in}}$
• Output: $Wx \in \mathbb{R}^{d_{ff} \times 1}$

GeGLU Dimensional Analysis

Let's trace through GeGLU with row vectors:

Input Space

• $x \in \mathbb{R}^{1 \times d_{model}}$ (e.g., $d_{model} = 768$ in the paper)

First Linear Transformations

• $W \in \mathbb{R}^{d_{model} \times d_{ff}}$
• $V \in \mathbb{R}^{d_{model} \times d_{ff}}$
• $b, c \in \mathbb{R}^{1 \times d_{ff}}$

Note: For GeGLU variants, they reduce $d_{ff}$ to $\frac{2}{3}$ of the original to maintain parameter count.

Intermediate Computations

1. Gate path: $xW + b \in \mathbb{R}^{1 \times d_{ff}}$
2. GELU applied element-wise: $\text{GELU}(xW + b) \in \mathbb{R}^{1 \times d_{ff}}$
3. Value path: $xV + c \in \mathbb{R}^{1 \times d_{ff}}$
4. Element-wise product: $\text{GELU}(xW + b) \otimes (xV + c) \in \mathbb{R}^{1 \times d_{ff}}$

Final Projection

• $W_2 \in \mathbb{R}^{d_{ff} \times d_{model}}$
• Output: $[\text{GELU}(xW + b) \otimes (xV + c)]W_2 \in \mathbb{R}^{1 \times d_{model}}$

Why Row Vectors?

1. Batch Processing

With row vectors, batching is natural:

• Batch input: $X \in \mathbb{R}^{B \times d_{model}}$
• Computation: $XW$ works directly
• Each row is an independent example

2. Implementation Efficiency

Most deep learning frameworks (PyTorch, TensorFlow) default to:

# Natural batched operation
output = torch.matmul(input, weight) + bias
# Shape: [batch, d_in] @ [d_in, d_out] = [batch, d_out]

3. Transformer Convention

The Transformer paper and most implementations use row vectors:

• Easier to think about sequence dimension
• Input shape: [batch, sequence, features]

Complete FFN_GeGLU Flow

Input: x ∈ ℝ^{1×d_model}
       ↓
   ┌───┴───┐
   ↓       ↓
  xW+b    xV+c     (both ∈ ℝ^{1×d_ff})
   ↓       ↓
 GELU()    │
   ↓       ↓
   └───⊗───┘       (element-wise multiply)
       ↓
     (·)W₂         (project back to d_model)
       ↓
Output: ∈ ℝ^{1×d_model}

Dimensional Example

If $d_{model} = 768$ and $d_{ff} = 2048$ (reduced from 3072):

• $x$: $[1 \times 768]$
• $W, V$: $[768 \times 2048]$
• $b, c$: $[1 \times 2048]$
• Intermediate: $[1 \times 2048]$
• $W_2$: $[2048 \times 768]$
• Output: $[1 \times 768]$

The row vector convention is simply more convenient for modern deep learning implementations!