# Understanding the Architecture of Needle 2: A Technical Breakdown of the Custom Transformer

> Explore the Needle 2 architecture. Discover its custom transformer design, Engram memory, Hadamard MLP, and nn.scan in JAX/Flax. Understand this advanced neural network.

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

---

**Needle 2 implements a custom transformer-style neural network in JAX/Flax with novel components including Engram memory tables, Hadamard MLP layers, and layer-wise scanning via `nn.scan`, all orchestrated through the `SimpleAttentionNetwork` class in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).**

Needle 2 is an open-source language model built on a custom transformer architecture that prioritizes efficient inference and flexible finetuning. Developed in the `cactus-compute/needle` repository, this JAX/Flax-based implementation introduces several specialized components not found in standard transformer implementations. The architecture centers on efficient attention mechanisms, compressed memory tables, and optimized feed-forward networks designed for modern hardware acceleration.

## Core Components of the Needle 2 Architecture

The architecture of Needle 2 is defined primarily in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), which contains specialized classes that replace or augment standard transformer building blocks to achieve high performance with reduced computational overhead.

### Configuration and Normalization

The `TransformerConfig` dataclass (lines 58-83) centralizes all hyperparameters including model dimensions, layer counts, attention heads, KV-window specifications, and Engram settings. This configuration enables easy swapping between model presets and ensures type-safe initialization across the stack.

For normalization, Needle 2 uses `ZCRMSNorm` (lines 46-55), a layer normalization variant that computes root-mean-square normalization with per-dimension learned scales. This approach provides numerical stability when operating in low-precision regimes such as bfloat16, which is critical for efficient training and inference.

### Attention and Memory Mechanisms

The `MultiHeadAttention` class implements standard multi-head self-attention with rotary position embeddings (RoPE) and Flash Attention support for GPU acceleration. A unique addition is the `Engram` class (lines 81-107), which functions as a learned memory table storing compressed token n-grams. Unlike standard KV caches that grow linearly with sequence length, the Engram provides additional key/value pairs through compressed representations, enabling longer context windows while respecting strict memory budgets defined by `KV_BUDGET_BYTES`.

### Feed-Forward Networks with Hadamard Transforms

Replacing the traditional two-linear feed-forward network, the `HadamardMLP` class (lines 80-103) implements a fast feed-forward layer using the Walsh-Hadamard transform. This architecture reduces computational complexity while maintaining non-linear capacity, offering a linear-time alternative to standard MLPs that typically require expensive matrix multiplications with expanded intermediate dimensions.

### Layer-wise Scanning and Block Stacking

Individual transformer layers are encapsulated in the `Block` class (lines 105-138), which stitches together attention, gating mechanisms (`attn_gate`), and the Hadamard MLP. Rather than using a Python loop over layers, Needle 2 employs `_ScanBody` (lines 140-176) and `Stack` (lines 176-227) classes to implement layer-wise scanning via `nn.scan`. This technique compiles the entire stack of `num_layers` blocks into a single JAX computation, dramatically reducing JIT compilation overhead and enabling automatic rematerialization (`cfg.remat`) for memory efficiency in deep models.

### Auxiliary Prediction Heads

Beyond standard language modeling, the architecture includes `ContrastiveHead` (lines 44-61) and `ConfidenceHead` (lines 63-75) classes for specialized downstream tasks. These heads pool hidden states using learned probes and project them into a lower-dimensional `contrastive_dim` space, enabling contrastive representation learning and confidence scoring for retrieval or tool-calling applications.

## Data Flow Through the Needle 2 Architecture

The forward pass through `SimpleAttentionNetwork` (lines 78-87) follows a distinct pipeline optimized for efficiency:

1. **Token Embedding**: Input tokens pass through `nn.Embed` scaled by √d to produce initial hidden states.
2. **RoPE Integration**: Rotary position embeddings are pre-computed via `precompute_rope_freqs` (line 97) and applied within the attention mechanism.
3. **Engram Injection**: The `Engram` generates compressed KV tables from token streams (`engram_indices`) and mixes these into the attention keys and values.
4. **Layer Scanning**: The `Stack` scans over `num_layers` `Block` instances, each executing:
   - RMS-normalized self-attention (`MultiHeadAttention`)
   - Gated residual connections (`attn_gate`)
   - Normalized Hadamard MLP (`HadamardMLP`)
5. **Aggregation**: The lane dimension is collapsed via mean aggregation, followed by final `ZCRMSNorm` normalization.
6. **Output Projection**: Logits are computed via dot-product between final hidden states and the transposed token embedding matrix.
7. **Auxiliary Outputs**: Optional contrastive embeddings and confidence scores are computed via their respective heads.

## Optimization Techniques in the Needle 2 Architecture

Several architectural decisions distinguish Needle 2 from standard transformer implementations:

- **Layer-wise Scanning**: The use of `nn.scan` in the `Stack` class reduces compilation time and enables efficient execution of deep models without unrolling the computation graph.
- **Engram Memory**: Compressed n-gram tables provide long-range context without linear memory growth, enforcing strict KV-budget constraints critical for inference optimization.
- **Hadamard Transform**: The `HadamardMLP` replaces dense matrix multiplications with fast Walsh-Hadamard transforms, reducing FLOPs while preserving model capacity.
- **Quantization Support**: Hooks like `_aq` and `maybe_quant_kv` enable 8-bit weight and activation quantization for accelerated inference.
- **Flash Attention**: Automatic utilization of optimized GPU kernels when available, falling back to manual attention implementations otherwise.

## Working with Needle 2: Practical Code Examples

The repository provides straightforward APIs for common tasks in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) and the CLI interface in [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py).

### Loading Checkpoints and Generating Text

```python
from needle.model.run import load_checkpoint, generate
from needle.model.architecture import SimpleAttentionNetwork
from needle.model.tokenizer import get_tokenizer

# Initialize model and load parameters

params, config = load_checkpoint("checkpoints/needle_step_1000.pkl")
model = SimpleAttentionNetwork(config)
tokenizer = get_tokenizer(config.vocab_size)

# Generate text

prompt = "The architecture of efficient transformers"
output = generate(model, params, tokenizer, prompt, max_new_tokens=64)
print(output)

```

### Extracting Contrastive Embeddings

```python
import jax.numpy as jnp

# Prepare input

tokens = tokenizer.encode("Example query for retrieval")
tokens = jnp.array([tokens])  # Shape: (1, seq_len)

model = SimpleAttentionNetwork(config)

# Get contrastive representation for downstream tasks

query_emb, _ = model.forward_contrastive(tokens, tokens)
print(query_emb.shape)  # (1, contrastive_dim)

```

### Accessing Hidden States for Interpretability

```python

# Extract internal representations for analysis

hidden = model.hidden_states(tokens)
print(hidden.shape)  # (num_layers, batch, seq, d_model)

# Example output: (27, 1, seq_len, 768)

```

### Running Inference via Command Line

```bash
needle run --checkpoint checkpoints/needle_step_1000.pkl \
           --query "Explain the architecture of Needle 2" \
           --max-len 128

```

## Summary

- Needle 2 is built in JAX/Flax with custom components defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), centering on the `SimpleAttentionNetwork` orchestrator.
- The architecture replaces standard components with optimized variants: `ZCRMSNorm` for stability, `HadamardMLP` for efficient feed-forward computation, and `Engram` for compressed memory.
- Layer-wise scanning via `nn.scan` in the `Stack` class enables efficient execution of deep transformer stacks without excessive compilation overhead.
- Built-in auxiliary heads (`ContrastiveHead`, `ConfidenceHead`) support retrieval and confidence estimation tasks beyond standard autoregressive modeling.

## Frequently Asked Questions

### What makes Needle 2's Engram memory different from standard transformer KV caches?

The `Engram` class (defined at lines 81-107 of [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)) stores compressed token n-grams as learned memory tables rather than maintaining full key-value pairs for every token in the sequence. This approach allows Needle 2 to respect a fixed `KV_BUDGET_BYTES` while still accessing long-range contextual information, effectively breaking the linear relationship between sequence length and memory consumption that plagues standard attention mechanisms.

### How does the Hadamard MLP improve efficiency over traditional feed-forward networks?

The `HadamardMLP` (lines 80-103 of [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)) replaces the standard two-layer feed-forward network—which requires expensive matrix multiplications with an expanded intermediate dimension—with a Walsh-Hadamard transform. This transform operates in linear time relative to model dimension, drastically reducing FLOPs while maintaining the non-linear capacity necessary for effective representation learning.

### Why does Needle 2 use layer-wise scanning (`nn.scan`) instead of a standard Python loop?

According to the implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `Stack` class uses `nn.scan` to handle the `_ScanBody` across all transformer layers (lines 140-227) in a single compiled JAX computation. This technique reduces JIT compilation overhead, enables automatic rematerialization (`cfg.remat`) for memory efficiency, and prevents the computation graph from becoming unwieldy when models contain approximately 27 layers or more.

### Where are the configuration options and model presets defined?

All architectural hyperparameters—including model size, attention head count, Engram settings, and KV-window specifications—are centralized in the `TransformerConfig` dataclass located at lines 58-83 of [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). This configuration object enables easy swapping between model presets and ensures consistent initialization of components like `MultiHeadAttention` and `Stack`.