Recurrent Memory › Implementing Transformer-XL

import numpy as np


class TransformerXLLayer:
    """
    A single Transformer-XL layer with segment-level recurrence
    and relative position encoding.
    """

    def __init__(
        self, d_model: int, n_heads: int, d_ff: int, max_rel_dist: int = 512
    ):
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.max_rel_dist = max_rel_dist

        # Initialize parameters (simplified: using random initialization)
        rng = np.random.default_rng(42)
        scale = 0.02

        # Attention projections
        self.W_q = rng.standard_normal((d_model, d_model)) * scale
        self.W_k_E = rng.standard_normal((d_model, d_model)) * scale  # Content
        self.W_k_R = rng.standard_normal((d_model, d_model)) * scale  # Position
        self.W_v = rng.standard_normal((d_model, d_model)) * scale
        self.W_o = rng.standard_normal((d_model, d_model)) * scale

        # Global biases for relative attention
        self.u = rng.standard_normal(d_model) * scale
        self.v = rng.standard_normal(d_model) * scale

        # Feed-forward network
        self.W_ff1 = rng.standard_normal((d_model, d_ff)) * scale
        self.b_ff1 = np.zeros(d_ff)
        self.W_ff2 = rng.standard_normal((d_ff, d_model)) * scale
        self.b_ff2 = np.zeros(d_model)

        # Layer norms (simplified: just storing means and variances)
        self.ln1_gamma = np.ones(d_model)
        self.ln1_beta = np.zeros(d_model)
        self.ln2_gamma = np.ones(d_model)
        self.ln2_beta = np.zeros(d_model)

        # Precompute relative position encodings
        self.rel_positions, self.rel_encodings = sinusoidal_relative_encoding(
            max_rel_dist, d_model
        )

    def layer_norm(
        self, x: np.ndarray, gamma: np.ndarray, beta: np.ndarray
    ) -> np.ndarray:
        """Apply layer normalization."""
        mean = x.mean(axis=-1, keepdims=True)
        std = x.std(axis=-1, keepdims=True) + 1e-6
        return gamma * (x - mean) / std + beta

    def relative_attention(
        self,
        hidden: np.ndarray,
        memory: np.ndarray,
    ) -> np.ndarray:
        """
        Compute relative multi-head attention with cached memory.

        Args:
            hidden: Current segment hidden states, shape (L, d_model)
            memory: Cached previous segment states, shape (M, d_model)

        Returns:
            Attention output, shape (L, d_model)
        """
        L = hidden.shape[0]
        M = memory.shape[0] if memory is not None else 0

        # Concatenate memory and hidden for keys/values
        if memory is not None and M > 0:
            extended = np.concatenate([memory, hidden], axis=0)
        else:
            extended = hidden
        K = extended.shape[0]

        # Compute queries, keys, values
        queries = hidden @ self.W_q  # (L, d)
        keys_E = extended @ self.W_k_E  # (K, d)
        values = extended @ self.W_v  # (K, d)

        # Compute attention scores with relative positions
        # Term (a): content-to-content
        scores = queries @ keys_E.T  # (L, K)

        # Terms (b), (c), (d): relative position terms
        # Simplified: we add positional bias based on distance
        for i in range(L):
            for j in range(K):
                # Compute relative distance
                if j < M:
                    key_pos = j - M  # Negative for memory
                else:
                    key_pos = j - M  # 0 to L-1 for current
                query_pos = i
                distance = query_pos - key_pos

                # Find relative encoding
                idx = np.where(
                    self.rel_positions
                    == np.clip(distance, -self.max_rel_dist, self.max_rel_dist)
                )[0]
                if len(idx) > 0:
                    r_ij = self.rel_encodings[idx[0]]
                    # Add position bias terms
                    scores[i, j] += (queries[i] + self.u) @ self.W_k_R @ r_ij
                    scores[i, j] += self.v @ self.W_k_R @ r_ij

        # Scale
        scores = scores / np.sqrt(self.d_head)

        # Causal mask (only for current segment attending to past)
        mask = np.ones_like(scores) * float("-inf")
        for i in range(L):
            # Can attend to all memory + tokens up to and including position i
            mask[i, : M + i + 1] = 0
        scores = scores + mask

        # Softmax
        exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
        attention_weights = exp_scores / (
            exp_scores.sum(axis=-1, keepdims=True) + 1e-8
        )

        # Output
        output = attention_weights @ values  # (L, d)
        output = output @ self.W_o

        return output, attention_weights

    def feed_forward(self, x: np.ndarray) -> np.ndarray:
        """Apply position-wise feed-forward network."""
        hidden = np.maximum(0, x @ self.W_ff1 + self.b_ff1)  # ReLU
        return hidden @ self.W_ff2 + self.b_ff2

    def forward(
        self, hidden: np.ndarray, memory: np.ndarray
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Forward pass for one layer.

        Args:
            hidden: Current segment hidden states
            memory: Cached memory from previous segment

        Returns:
            output: New hidden states
            new_memory: States to cache for next segment
        """
        # Self-attention with memory
        attn_out, attn_weights = self.relative_attention(hidden, memory)
        hidden = self.layer_norm(
            hidden + attn_out, self.ln1_gamma, self.ln1_beta
        )

        # Feed-forward
        ff_out = self.feed_forward(hidden)
        output = self.layer_norm(hidden + ff_out, self.ln2_gamma, self.ln2_beta)

        return output, attn_weights

# Demonstrate the layer
layer = TransformerXLLayer(d_model=64, n_heads=4, d_ff=256, max_rel_dist=32)

# Simulate processing two segments
segment_length = 8
hidden_dim = 64

# Segment 1: no memory yet
rng = np.random.default_rng(123)
segment1_input = rng.standard_normal((segment_length, hidden_dim)) * 0.5
segment1_output, attn1 = layer.forward(segment1_input, memory=None)

# Segment 2: use segment 1's output as memory
segment2_input = rng.standard_normal((segment_length, hidden_dim)) * 0.5
segment2_output, attn2 = layer.forward(segment2_input, memory=segment1_output)