# Understanding Sandwich Normalization in Needle 2: Architecture and Implementation

> Discover Sandwich Normalization in Needle 2 a residual connection technique that stabilizes gradients and enables dynamic information routing. Learn its architecture and implementation.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: deep-dive
- Published: 2026-09-08

---

**Sandwich Normalization in Needle 2 is a residual connection technique that applies Root Mean Square (RMS) normalization before and after each transformer sub-layer, then fuses the results through a learned sigmoid gate to stabilize gradients and enable dynamic information routing.**

Sandwich Normalization is a core architectural innovation in the Needle 2 transformer model developed by Cactus Compute. According to the cactus-compute/needle source code, this technique replaces traditional layer normalization with a dual RMSNorm approach wrapped around attention and MLP blocks, creating residual connections that adaptively control information flow through learned gating parameters.

## What Is Sandwich Normalization?

Sandwich Normalization is a residual connection pattern that "sandwiches" each transformer sub-layer between two normalization operations. Unlike standard Pre-LN or Post-LN architectures, this approach normalizes both the input to the sub-layer and the output from the sub-layer before combining them.

Mathematically, for an input $x$ and sub-layer function $F$ (attention or MLP):

```python
h1 = RMSNorm(x)              # First slice (pre-normalization)

h2 = RMSNorm(F(h1))          # Second slice (post-normalization)

gate = sigmoid(g1 * h1 + g2 * h2 + b)  # Learned per-channel gate

output = x + gate * h2       # Sandwich-normed residual

```

This creates a "sandwich" structure where the computational meat of the transformer block sits between two RMS normalization layers.

## How Sandwich Normalization Works in Needle 2

### The Architecture Pipeline

In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), each `TransformerBlock` implements the sandwich pattern through a strict data flow:

1. **Input Residual** passes through the first **RMSNorm** layer
2. The normalized tensor feeds into the core sub-layer (multi-head attention or MLP)
3. The sub-layer output undergoes a second **RMSNorm** transformation
4. A **learned gate** combines the pre- and post-normalized tensors using a sigmoid activation
5. The gated result is added back to the original residual via a skip connection

### RMS Normalization Details

**RMS Normalization** divides each vector by its root-mean-square value, maintaining signal scale without introducing learnable affine parameters (gamma/beta). As implemented in the Needle 2 architecture, this element-wise operation adds minimal computational overhead while preventing parameter drift during fine-tuning.

### The Learned Gating Mechanism

The gating mechanism uses a **sigmoid-activated** linear combination of the two normalized states:

```

gate = σ(g₁·h₁ + g₂·h₂ + b)

```

Here, $g_1$, $g_2$, and $b$ are learned per-channel parameters that allow the network to dynamically decide how much of the transformed signal to merge with the original context. This is particularly crucial for Needle 2's mixed-type tool-calling tasks, where the model must selectively preserve or replace information.

## Implementation in the Source Code

The concrete implementation resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). The `TransformerBlock` class encapsulates the sandwich logic within its forward pass, utilizing `rms_norm` methods applied sequentially to ensure the dual-normalization pattern.

```python
from needle.model.architecture import TransformerBlock

# Initialize a transformer block with sandwich normalization

block = TransformerBlock(hidden_dim=512, n_heads=8)

# The internal forward pass executes:

# 1. x_norm = block.rms_norm(x)           # First RMSNorm slice

# 2. att_out = block.attention(x_norm)    # Attention sub-layer

# 3. att_norm = block.rms_norm(att_out)   # Second RMSNorm slice

# 4. gated = block.gate(x_norm, att_norm) # Learned sigmoid gating

# 5. y = x + gated * att_norm             # Sandwich-normed residual

```

## Key Benefits of Sandwich Normalization

- **Training Stability**: By removing learnable scale and shift parameters from the normalization layers, RMSNorm prevents drift during fine-tuning, making the architecture robust for downstream adaptation.
- **Computational Efficiency**: RMSNorm operates element-wise without additional parameters, adding negligible overhead compared to standard LayerNorm.
- **Expressive Control**: The dual normalization combined with learned gating allows dynamic routing of information, enabling the model to balance retention of original context against integration of new transformed features.

## Practical Usage Examples

### Basic Model Instantiation

When loading Needle 2, sandwich normalization is automatically applied within every transformer block:

```python
import needle

# Load the pre-packed 14MB engine

agent = needle.Needle(weights="cactus-needle/needle2.cact")

# Sandwich-normed residuals are active in all attention and MLP blocks

```

### Fine-Tuning with LoRA Adapters

The normalization layers remain frozen during adapter-based fine-tuning, preserving the sandwich architecture's stability:

```python
import needle

# Load base model with LoRA adapter

agent = needle.Needle(
    weights="cactus-needle/needle2.cact",
    lora="my_lora_adapter.pkl"
)

# Execute tool-calling; sandwich-norm parameters remain unchanged

result = agent.run("log lunch, turkey sandwich on rye, and 600ml of water")
print(result["results"])

# [{'description': 'turkey sandwich on rye', 'meal_type': 'lunch'},

#  {'amount_ml': 600}]

```

## Key Source Files

- **[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)**: Contains the `TransformerBlock` implementation with RMS normalization and gating logic for sandwich-normed residuals.
- **[`README.md`](https://github.com/cactus-compute/needle/blob/main/README.md)**: Documents the "Simple Attention Network" design rationale and the sandwich normalization concept.
- **[`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py)**: Demonstrates that sandwich normalization logic persists through model quantization.
- **[`needle/_worker.py`](https://github.com/cactus-compute/needle/blob/main/needle/_worker.py)**: Entry point connecting the sandwich architecture to the public `Needle` runtime API.

## Summary

- Sandwich Normalization wraps each transformer sub-layer with dual RMS normalization stages, creating a "sandwich" structure that normalizes both inputs and outputs.
- The architecture uses element-wise RMSNorm to maintain stable signal scales without learnable affine parameters.
- A learned sigmoid gate dynamically fuses pre- and post-normalized tensors, enabling adaptive residual connections.
- Implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), this pattern is critical for Needle 2's tool-calling capabilities and training stability.
- The normalization layers remain fixed during LoRA fine-tuning, preventing drift while allowing task-specific adaptation.

## Frequently Asked Questions

### How does Sandwich Normalization differ from Pre-LN and Post-LN?

Pre-LN applies normalization only before the sub-layer, while Post-LN applies it after. Sandwich Normalization applies RMS normalization both before and after, then uses a learned gate to combine these states. This dual normalization provides more stable gradients than Post-LN and better expressive control than Pre-LN.

### Why does Needle 2 use RMSNorm instead of LayerNorm?

RMSNorm eliminates the learnable gain and bias parameters found in LayerNorm, dividing inputs solely by their root-mean-square. This reduces parameter count and prevents normalization drift during fine-tuning, which is essential for efficient adapter-based training in the 14MB Needle 2 engine.

### Where is the gating mechanism implemented in the source code?

The gating mechanism is implemented within the `TransformerBlock` class in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). It computes a sigmoid-activated linear combination of the pre-normalized and post-normalized tensors using learned per-channel weights, then applies this gate to the residual connection.

### Does quantization affect the sandwich normalization layers?

No. As shown in [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py), the sandwich normalization logic remains intact after quantization. The RMS normalization operations and gating calculations persist in the quantized engine, ensuring consistent behavior between full-precision and quantized Needle 2 models.