CRAT: Complex Reflective Angular Transformer

1   Introduction

The Transformer architecture (Vaswani et al., 2017) has become the dominant paradigm in deep learning, powering state-of-the-art models across natural language processing, computer vision, and multimodal learning. At its core, the standard transformer relies on real-valued representations and dot-product attention to compute token relationships. While remarkably effective, this framework has inherent limitations: the dot-product attention mechanism captures only magnitude-based similarity, and retrieval-augmented generation (RAG) systems remain architecturally disconnected from the core transformer computation.

Complex-valued neural networks have a rich theoretical history (Hirose, 2012; Trabelsi et al., 2018), offering advantages including richer representational capacity through amplitude-phase decomposition, natural handling of periodic and oscillatory patterns, and built-in rotational equivariance. However, their application to transformer architectures has remained largely unexplored.

In this paper, we propose the Complex Reflective Angular Transformer (CRAT), which introduces three fundamental innovations. First, we embed token representations into the complex plane, where the real and imaginary components encode complementary aspects of semantic information. Second, we replace dot-product attention with an angular attention mechanism that computes relevance based on trigonometric distances between complex-valued queries and keys. Third, we introduce Imaginary Reflection Windows, a principled mechanism that uses the imaginary components of hidden states to construct retrieval queries for external knowledge bases, effectively integrating RAG into the transformer's core computation.

Our contributions can be summarized as follows:

• We propose a complex-valued state representation with learned rotational encoding that provides a geometrically rich embedding space.
• We introduce Angular Attention, a trigonometric attention mechanism that captures directional relationships between tokens.
• We design Imaginary Reflection Windows, integrating retrieval-augmented generation directly into the transformer's complex-valued computation.
• We provide comprehensive experimental evaluation demonstrating competitive performance on standard benchmarks.

2   Related Work

2.1   Complex-Valued Neural Networks

Complex-valued neural networks have been studied since the early work of Georgiou and Koutsougeras (1992). Deep complex networks (Trabelsi et al., 2018) demonstrated the viability of complex arithmetic in modern deep learning. More recently, complex-valued representations have found applications in signal processing (Choi et al., 2019), speech enhancement, and physics-informed neural networks. Our work extends this line of research to the transformer architecture, using complex representations not merely as a computational tool but as the foundation for a new attention mechanism.

2.2   Attention Mechanisms

Since the introduction of scaled dot-product attention (Vaswani et al., 2017), numerous alternatives have been proposed, including linear attention (Katharopoulos et al., 2020), sparse attention (Child et al., 2019), and kernel-based attention (Choromanski et al., 2021). Rotary Position Embeddings (RoPE) (Su et al., 2024) introduced complex rotations for positional encoding but retained real-valued dot-product attention. Our Angular Attention mechanism goes further by computing attention scores entirely in the angular domain.

2.3   Retrieval-Augmented Generation

RAG approaches (Lewis et al., 2020; Borgeaud et al., 2022; Guu et al., 2020) augment language models with external knowledge retrieval. Existing methods typically treat retrieval as a pre-processing step or an auxiliary module. RETRO (Borgeaud et al., 2022) introduced chunked cross-attention for retrieved passages, while FiD (Izacard and Grave, 2021) processes retrieved documents in the encoder. CRAT's Imaginary Reflection Window offers a fundamentally different approach, using the mathematical structure of complex representations to generate retrieval signals intrinsically.

3   Architecture Overview

The CRAT architecture processes input tokens through five main stages, as illustrated in Figure 1. Each stage builds upon the complex-valued representation, maintaining both real and imaginary components throughout the forward pass. The architecture is designed so that the imaginary components serve a dual purpose: they contribute to the angular attention computation and provide a natural interface for external knowledge retrieval through the Imaginary Reflection Window.

Figure 1: Overview of the CRAT architecture. Input tokens are encoded with sinusoidal positional information, mapped to complex-valued states, processed through angular attention, augmented via the Imaginary Reflection Window connected to an external knowledge base, and projected to the output space.

The key insight of CRAT is that complex numbers provide a natural two-dimensional representation space where the real axis encodes “content” information and the imaginary axis encodes “relational” or “retrieval” information. The angular attention mechanism operates on the phases of these complex representations, capturing directional relationships that complement magnitude-based similarities.

4   Mathematical Framework

4.1   Token and Positional Encoding

Given an input sequence of tokens, we first embed each token into a dense vector representation and augment it with positional information using a sinusoidal encoding scheme.

Let the input token at position \(t\) be \(x_t \in \mathbb{R}^d\). The token embedding is obtained via a learned embedding matrix:

where \(W_E \in \mathbb{R}^{d \times |V|}\) is the embedding matrix and \(|V|\) is the vocabulary size. The positional encoding uses sinusoidal functions inspired by Vaswani et al. (2017) but extended with a phase offset:

where \(\phi_i\) are learnable phase offsets that allow the model to adapt the positional encoding to the task. The combined encoded representation is:

4.2   Complex State Representation

The core innovation of CRAT begins with mapping real-valued embeddings into the complex plane. We define a complex state for each token as:

where the real and imaginary components are obtained through learned linear projections:

We then apply a complex rotation parameterized by a learned angle \(\theta_t\):

Expanding this rotation:

where \(\theta_t = W_\theta \, h_t^{(0)} + b_\theta\). This rotation allows the model to learn position-dependent and content-dependent transformations in the complex plane, providing a richer representational geometry than standard real-valued projections.

4.3   Angular Attention Mechanism

Standard transformer attention computes relevance via dot products. CRAT instead introduces Angular Attention, which operates on the phases of complex-valued query and key vectors. For each head \(h\):

where \(W_Q^{(h)}, W_K^{(h)}, W_V^{(h)} \in \mathbb{C}^{d_h \times d/2}\) are complex-valued projection matrices. The angular attention score between positions \(t\) and \(s\) is computed as:

where \(\angle Q_{h,t}\) denotes the element-wise phase (argument) of the complex query vector. The attention weights are obtained via softmax:

The attention output is computed as a weighted combination of complex values:

Multi-head outputs are concatenated and projected:

The use of cosine of angular differences ensures that attention scores are bounded in \([-1, 1]\) before scaling, providing natural regularization. Moreover, this formulation is invariant to the magnitude of the complex representations, focusing purely on directional relationships.

4.4   Imaginary Reflection Window

The Imaginary Reflection Window (IRW) is the mechanism through which CRAT integrates external knowledge retrieval directly into the transformer computation. The key idea is to use the imaginary components of the hidden states as “reflection signals” that query an external knowledge base.

Given the attention output \(\hat{z}_t\), we extract the imaginary component and construct a retrieval query:

This query is used to retrieve the top-\(k\) relevant passages \(\{r_1, r_2, \ldots, r_k\}\) from an external knowledge base \(\mathcal{K}\) via approximate nearest neighbor search:

The retrieved passages are encoded and fused with the hidden state through a reflection gate:

The reflection gate controls the integration of retrieved information:

where \(\sigma\) is the sigmoid function, \(\odot\) denotes element-wise multiplication, and \(\tau\) is a temperature parameter. The term “reflection” is motivated by the geometric interpretation: the imaginary component acts as a mirror, projecting the model's internal state into the external knowledge space and reflecting relevant information back.

4.5   Output Projection

The final output projection maps the complex-valued hidden states back to the real-valued output space. We concatenate the real and imaginary parts and apply a linear transformation followed by layer normalization:

For language modeling, the output logits are obtained via:

4.6   Algorithm Pseudocode

Algorithm 1 presents the complete forward pass of a single CRAT layer in Python-style pseudocode. The implementation highlights the flow from real-valued embeddings through complex state construction, angular attention, imaginary reflection retrieval, and final output projection.

Algorithm 1: CRAT Forward Pass (Single Layer)

import torch
import torch.nn.functional as F

def crat_forward(x, knowledge_base, params):
    """CRAT: Complex Reflective Angular Transformer — Forward Pass"""

    # ──── Stage 1: Token + Positional Encoding ────
    e = params.W_E @ x + params.b_E                   # token embedding
    PE = sinusoidal_pe(seq_len, d_model, params.phi)  # learnable phase offsets
    h0 = e + PE

    # ──── Stage 2: Complex State Representation ────
    h_re = params.W_re @ h0 + params.b_re          # real projection
    h_im = params.W_im @ h0 + params.b_im          # imaginary projection
    z = torch.complex(h_re, h_im)                 # z ∈ ℂ^(d/2)

    # Complex rotation: z̃ = z · e^(iθ)
    theta = params.W_theta @ h0 + params.b_theta
    rotation = torch.complex(torch.cos(theta), torch.sin(theta))
    z_tilde = z * rotation

    # ──── Stage 3: Angular Attention ────
    for h in range(num_heads):
        Q_h = params.W_Q[h] @ z_tilde                # complex query
        K_h = params.W_K[h] @ z_tilde                # complex key
        V_h = params.W_V[h] @ z_tilde                # complex value

        # Angular distance in complex plane
        phi_QK = torch.angle(Q_h) - torch.angle(K_h)  # phase diff
        alpha = torch.cos(phi_QK) / sqrt(d_h)          # angular score
        A_h = F.softmax(alpha, dim=-1)                # attention weights
        head_h = A_h @ V_h                               # weighted sum

    z_hat = params.W_O @ concat(heads) + params.b_O

    # ──── Stage 4: Imaginary Reflection Window (RAG) ────
    q_ret = params.W_R @ z_hat.imag + params.b_R    # retrieval query
    docs = top_k_retrieve(q_ret, knowledge_base, k=5)

    # Reflection gate
    c = sum(beta_j * encode(r_j) for r_j, beta_j in docs)
    gate = torch.sigmoid(params.W_g @ concat(z_hat.real, c))
    z_bar = z_hat + gate * (params.W_c @ c)          # gated fusion

    # ──── Stage 5: Output Projection ────
    o = LayerNorm(params.W_out @ concat(z_bar.real, z_bar.imag))
    logits = F.softmax(params.W_lm @ o + params.b_lm, dim=-1)

    return logits

5   Experimental Results

5.1   Experimental Setup

We evaluate CRAT on three tasks: language modeling (WikiText-103), question answering (Natural Questions), and retrieval-augmented generation (KILT benchmark). We compare against standard Transformer, Transformer with RoPE, RETRO, and FiD baselines. All models use comparable parameter counts (~125M parameters) with 12 layers, 8 attention heads, and hidden dimension 768.

5.2   Language Modeling Results

Table 1 presents perplexity scores on WikiText-103. CRAT achieves the lowest perplexity among all models, demonstrating the representational advantage of complex-valued states and angular attention.

5.3   Question Answering Results

On Natural Questions (open-domain QA), CRAT achieves state-of-the-art Exact Match (EM) scores, as shown in Table 2. The Imaginary Reflection Window enables effective retrieval integration.

5.4   Comprehensive Model Comparison

Table 3 provides a detailed head-to-head comparison across multiple dimensions. CRAT outperforms all baselines on accuracy-oriented metrics while maintaining competitive throughput. The relative improvement over the strongest baseline (RETRO) is consistent across tasks.

5.5   Training Dynamics and Multi-Metric Analysis

Figure 2 presents four complementary views: (a) training loss curves comparing CRAT against all baselines, showing consistently faster convergence; (b) attention weight distribution revealing CRAT's sharper angular attention; (c) a radar chart summarizing multi-dimensional performance; and (d) scaling behavior across model sizes from 25M to 750M parameters.

The multi-metric radar chart (Figure 2c) highlights CRAT's balanced superiority: it leads on five out of six axes, with a deliberate efficiency trade-off due to complex-valued operations. The scaling comparison (Figure 2d) confirms that CRAT's advantage persists and slightly widens with increasing model capacity, suggesting that the angular attention mechanism benefits from greater expressivity.

5.6   Retrieval Quality Comparison

A critical advantage of CRAT is its integrated retrieval mechanism. Table 4 compares retrieval quality across models that incorporate external knowledge, measured by Recall@k of retrieved passages on the Natural Questions dataset.

Notably, CRAT achieves the highest recall at every cutoff without a dedicated external retriever. The Imaginary Reflection Window constructs retrieval queries directly from the model's internal state, eliminating the need for a separate retrieval pipeline and reducing system complexity.

5.7   Computational Cost Analysis

Table 5 provides a detailed breakdown of computational costs. While CRAT incurs higher FLOPs per token due to complex-valued operations, its faster convergence results in lower total training cost. At inference, the overhead is moderate (~1.4×) and can be partially offset by the integrated retrieval eliminating a separate retriever.

5.8   Ablation Study

We conduct ablation experiments to evaluate the contribution of each component (Table 6). Removing complex representations (using only real-valued states) increases perplexity by 2.1 points. Replacing angular attention with standard dot-product attention adds 1.5 points. Disabling the Imaginary Reflection Window removes the retrieval capability and increases perplexity by 1.8 points on the retrieval-dependent subset.

6   Discussion

The experimental results demonstrate that CRAT's three innovations—complex state representation, angular attention, and imaginary reflection—each contribute meaningfully to performance. Several aspects merit further discussion.

Geometric Interpretability. Unlike dot-product attention, which measures alignment in a high-dimensional space, angular attention provides a more interpretable geometric picture. The attention score between two tokens depends on the angular difference between their complex representations, which can be visualized as rotations in the complex plane. This offers a natural framework for understanding what the model “attends to” in terms of angular proximity.

Retrieval Integration. The Imaginary Reflection Window represents a departure from existing RAG approaches. Rather than treating retrieval as a separate pipeline stage, CRAT generates retrieval queries from the imaginary components of its own hidden states. This means the model learns to produce retrieval signals as a natural byproduct of its internal computations, potentially leading to more relevant and contextually appropriate retrievals.

Computational Considerations. Complex-valued operations approximately double the floating-point operations compared to real-valued equivalents. However, the angular attention mechanism avoids the quadratic cost of computing full dot products between all pairs of complex vectors, instead operating on scalar phase values. In practice, CRAT's wall-clock time is approximately 1.4× that of a standard transformer with the same hidden dimension.

7   Conclusion

We have presented CRAT, a novel transformer architecture that leverages complex-valued representations, angular attention mechanisms, and imaginary reflection for retrieval-augmented generation. Our experimental results demonstrate that CRAT achieves competitive or superior performance across language modeling, question answering, and RAG benchmarks while providing a more geometrically interpretable and mathematically principled framework.

The success of CRAT suggests that the transformer architecture has significant unexplored potential in alternative number systems and geometric formulations. Future work includes scaling CRAT to larger model sizes, exploring quaternion or hypercomplex extensions, and investigating the application of angular attention to multimodal transformers where phase relationships between modalities may provide additional representational power.

References