# Simple Attention Network Architecture Explained: Core Components and Implementation in Needle

> Explore the Simple Attention Network architecture including multi-head attention, engram memory, and dual output heads. Learn its core components and implementation in Needle.

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

---

**The Simple Attention Network is a transformer-style encoder that combines multi-head attention with auxiliary components including engram memory, multi-token prediction, and dual output heads for contrastive learning and confidence scoring.**

The `SimpleAttentionNetwork` class in the `cactus-compute/needle` repository implements a modular, research-friendly architecture designed for efficient sequence modeling and self-supervised learning. This article breaks down its implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), explaining how each component contributes to the forward pass and what makes this architecture distinctive.

## Core Components of the Simple Attention Network

The network is constructed in the `setup` method and comprises eight main sub-modules. Each serves a specific purpose in the processing pipeline.

### Embedding Layer and Scaling

The first operation transforms discrete tokens into continuous representations.

In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) lines 83-84, the embedding layer and scaling factor are initialized:

```python
self.embedding = nn.Embed(cfg.vocab_size, cfg.d_model, dtype=cfg.jax_dtype)
self.embed_scale = jnp.sqrt(cfg.d_model).astype(cfg.jax_dtype)

```

Token embeddings are scaled by √d_model to stabilize gradients during the early training stages. This follows the original Transformer convention with a minor adaptation for JAX dtype handling.

### Transformer Stack

The backbone of the network is a deep stack of transformer blocks.

Instantiated at line 85 as `self.stack = Stack(cfg)`, this module processes the embedded sequence through multiple layers of multi-head attention and feed-forward computation. The `Stack` class (defined elsewhere in the same file) uses `HadamardMLP` for the feed-forward sub-layers and supports efficient attention variants.

### Dual Output Heads: Contrastive and Confidence

The architecture produces multiple output streams beyond standard language modeling logits.

**Contrastive head** (lines 86-88):

```python
self.contrastive_head = ContrastiveHead(cfg.d_model, cfg.contrastive_dim, cfg.jax_dtype)

```

Projects hidden states to a lower-dimensional space for self-supervised similarity learning.

**Confidence head** (lines 89-90):

```python
self.confidence_head = ConfidenceHead(cfg.jax_dtype)

```

Predicts a scalar confidence score per token position, useful for uncertainty quantification and selective prediction.

### Engram Memory: Learned N-gram Statistics

A distinctive feature of this architecture is the **engram memory** — learned lookup tables that augment attention with statistical n-gram patterns.

Constructed in lines 92-95 via list comprehension:

```python
self.engram_k = [nn.Embed(cfg.engram_vocab, cfg.d_kv, dtype=cfg.jax_dtype)
                 for _ in range(cfg.num_engram_layers)]
self.engram_v = [nn.Embed(cfg.engram_vocab, cfg.d_kv, dtype=cfg.jax_dtype)
                 for _ in range(cfg.num_engram_layers)]

```

These tables provide additional key/value pairs to selected attention layers, effectively injecting prior distributional knowledge without increasing the main sequence length. The `_engram_kv` method fetches these values during the forward pass.

### Multi-Token Prediction Pathway

The **MTP (Multi-Token Prediction)** pathway enables training on multiple future tokens simultaneously, improving sample efficiency.

Created in lines 96-102:

```python
self.mtp_combine = nn.Dense(cfg.d_model, dtype=cfg.jax_dtype)
self.mtp_block = Block(cfg, layer_idx=cfg.num_layers, is_mtp=True)
self.mtp_emb_norm = nn.RMSNorm(dtype=cfg.jax_dtype)
self.mtp_final_norm = nn.RMSNorm(dtype=cfg.jax_dtype)

```

This pathway combines the final stack output with the next-token embedding, processes through an additional block, and produces auxiliary logits.

### Rotary Positional Encoding

Position information is injected via **RoPE (Rotary Position Embedding)**.

The `_rope` method (lines 104-108) pre-computes frequency tensors:

```python
def _rope(self):
    inv_freq = 1.0 / (cfg.rope_theta ** (jnp.arange(0, cfg.d_head, 2).astype(cfg.jax_dtype) / cfg.d_head))
    # ... returns cosine/sine precomputed for all positions up to max_seq_len

```

RoPE enables better length generalization than absolute positional embeddings by encoding relative positions through rotation matrices.

## Forward Pass Execution Flow

The `__call__` method orchestrates these components in a specific sequence:

1. **Embedding**: `x = self.embedding(tokens) * self.embed_scale`
2. **RoPE generation**: `rope = self._rope()`
3. **Engram lookup** (optional): `engram_k, engram_v = self._engram_kv(...)`
4. **Stack processing**: `x = self.stack(x, rope, engram_k, engram_v, mask=...)`
5. **Logit projection**: `logits = self.embedding.attend(x)` (weight tying)
6. **MTP computation** (if requested): combines `x` with next-token embeddings through the MTP pathway

The implementation supports optional masking utilities including `make_causal_mask` and `make_padding_mask` for controlling attention patterns.

## Practical Usage Examples

### Basic Initialization and Forward Pass

```python
import jax.numpy as jnp
from needle.model.architecture import SimpleAttentionNetwork, TransformerConfig

# Configuration with 12 layers, 6 KV heads (GQA), 2K context length

cfg = TransformerConfig(
    vocab_size=32000,
    d_model=768,
    num_layers=12,
    num_heads=12,
    num_kv_heads=6,        # Grouped-query attention

    max_seq_len=2048,
    contrastive_dim=256,    # For contrastive head

    num_engram_layers=4     # Enable engram on last 4 layers

)

model = SimpleAttentionNetwork(cfg)

```

### Training-Time Forward with All Outputs

```python
tokens = jnp.ones((2, 16), dtype=jnp.int32)  # batch=2, seq_len=16

# Full forward with MTP auxiliary loss

outputs = model(tokens, return_mtp=True, deterministic=False)

# outputs contains:

# - 'logits': (2, 16, 32000) — main LM predictions

# - 'mtp_logits': (2, 16, 32000) — auxiliary MTP predictions

```

### Inference-Time Auxiliary Head Usage

```python

# Extract contrastive embeddings for retrieval

contrastive_emb = model.encode_contrastive(tokens)

# shape: (2, 16, 256)

# Get per-token confidence scores

confidence = model.forward_confidence(tokens)

# shape: (2, 16) — higher values indicate more confident predictions

```

## Key Implementation Files

| File | Purpose |
|------|---------|
| [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) | Core `SimpleAttentionNetwork` class with all sub-modules |
| [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py) | Fake-quantization utilities (`maybe_quant_kv`) for efficient attention |
| [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) | Model initialization and inference wrappers |
| [`tests/test_inference.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_inference.py) | Validation suite for forward pass and auxiliary heads |

## Distinctive Design Decisions

The Simple Attention Network architecture in `cactus-compute/needle` differs from standard transformers in three key ways:

- **Engram memory** adds learned n-gram statistics as attention-augmenting KV pairs, bridging statistical and neural language modeling
- **Dual heads** enable self-supervised objectives (contrastive learning) and uncertainty estimation without architectural branching
- **MTP pathway** provides an auxiliary prediction task that improves data efficiency during pretraining

These components are modular — engrams can be disabled via config, and auxiliary heads are only computed when requested.

## Summary

- **SimpleAttentionNetwork** implements a transformer encoder with embeddings scaled by √d_model, deep `Stack` processing, and tied input/output embeddings
- **Engram memory** provides learned lookup tables that augment attention with n-gram statistics on selected layers
- **MTP pathway** enables multi-token prediction through a secondary processing branch combining final hidden states with next-token embeddings
- **Dual output heads** support contrastive representation learning (`ContrastiveHead`) and confidence scoring (`ConfidenceHead`)
- **RoPE positional encoding** handles position information via pre-computed rotary frequencies
- All components are configured through `TransformerConfig` and implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)

## Frequently Asked Questions

### What makes the Simple Attention Network different from a standard Transformer?

The key additions are engram memory for statistical augmentation, multi-token prediction for improved sample efficiency, and dual output heads for auxiliary self-supervised tasks. The core attention mechanism remains compatible with efficient implementations like FlashAttention.

### How does engram memory work in practice?

Engram memory consists of learned embedding tables (`engram_k` and `engram_v`) that provide additional key/value pairs to specified attention layers. During forward pass, the `_engram_kv` method fetches these based on token n-grams, augmenting the standard self-attention computation without increasing sequence length.

### When should I use the MTP pathway?

Enable MTP (`return_mtp=True`) during training to add an auxiliary prediction loss for the next token. This improves sample efficiency and can accelerate convergence. During standard inference, you can disable it for marginally faster computation.

### What is the confidence head used for?

The confidence head outputs a scalar per position indicating model certainty. This enables applications like selective prediction (abstaining on low-confidence outputs), uncertainty calibration, and confidence-based sampling strategies during generation.