💻 computer science

One Vector Is All You Need for O(1) Self-Attention: The Ocean State

This paper introduces the "Ocean State," a method that replaces standard self-attention with a single persistent vector to achieve O(1) computation and memory complexity while demonstrating superior training stability and no signs of forgetting even at 10 million steps of context.

Original authors: SHUYUAN YU

Published 2026-09-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: SHUYUAN YU

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Modern artificial intelligence often relies on a specific type of computer program called a Transformer, which has become the standard for tasks like writing text, translating languages, and answering questions. These programs work by looking at a sequence of words and deciding what word comes next. To do this accurately, the program must remember the words it has already seen. In the current design, as the program reads a longer sentence or a whole book, it keeps a growing list of every word it has processed so far. This list acts like a memory bank that expands with every new word. While this works well for short texts, it becomes a heavy burden for very long ones. The computer must constantly scan this ever-growing list to find relevant information, which requires a massive amount of processing power and memory. As the text gets longer, the time and energy needed to read it grow much faster than the text itself, eventually making it impossible to process very long documents efficiently.

A researcher named Yu Shuyuan has proposed a different way to handle this memory problem. Instead of keeping a growing list of every word, the new method suggests keeping just one single summary vector, which the author calls the "ocean state." Imagine this state as a single, dense container that holds the essence of everything the program has read so far. As the program reads a new word, it updates this single container to include the new information, replacing the old summary with a fresh one. The program then uses this single container to help predict the next word. This approach changes the fundamental math of the problem: instead of the effort growing with the length of the text, the effort stays the same no matter how long the text becomes. The program can read a million words with the same amount of work per word, and it uses a fixed, tiny amount of memory regardless of the length.

The core idea behind this work is a simple question: what if we did not throw away the final summary of a sentence after using it to predict the next word? In standard programs, this summary is calculated and then discarded. The new method keeps it, treating it as a persistent memory that flows from one step to the next. The program reads the current word and this single summary vector, combines them, and produces a new summary. This process repeats for every word in the text. Crucially, the program learns how to read and update this single vector entirely on its own. No one told it how to compress information into one spot; it simply learned to do so during training because the design allowed it to happen. The researchers found that this single vector is surprisingly capable. It can hold information from a sequence of ten million steps without losing any of it. In tests, the program could recall specific words from the beginning of a ten-million-step sequence just as accurately as it could from the first few steps, showing no sign of forgetting.

The researchers tested this idea against the standard method using a large set of text data. They trained both versions to predict the next word in a sentence. The results showed that the new method, which uses only one vector, consistently performed better than the standard method that keeps a growing list. This advantage was not a fluke; it appeared across many different settings, including different sizes of models and different lengths of text. In fact, the new method was so stable that it could process a sequence of ten million steps without any increase in error, whereas the standard method would have run out of memory or taken an impossible amount of time to process even a fraction of that length. The researchers also tested a specific task where the program had to repeat a word after a long delay. The new method reproduced the word with a loss of 0.0006, while the standard method failed completely.

One concern with such a compressed memory is that it might become unreadable or "foggy" over time, losing the details needed to make good predictions. However, the experiments showed the opposite. The single vector remained sharp and clear, capable of retrieving specific information even after millions of steps. The researchers also explored how to make this method faster on real computers. Because the new method processes words one after another in a strict sequence, it can be slower to train on a single machine compared to the standard method, which can process many words at once. To solve this, the researchers developed a way to split the work across different layers of the program, allowing multiple steps to be processed simultaneously. This engineering change significantly reduced the time needed to train the model while keeping the same high performance.

The study demonstrates that it is possible to build a system that remembers a vast amount of information using a fixed amount of memory and processing power. The researchers showed that the program learns to organize this information efficiently without any special instructions. While the current experiments were conducted on a specific dataset and with relatively small models, and behavior at real-world sizes remains untested, the results suggest that the underlying principle is robust. The method works by repositioning how the program handles its own internal state, turning a discarded summary into a persistent memory. This change eliminates the structural cost that has limited the length of text AI systems can handle. The findings indicate that with this approach, the limitations of context length may no longer be a barrier, allowing for systems that can read and understand documents of any size with consistent speed and accuracy.

Technical Summary: One Vector Is All You Need for O(1) Self-Attention: The Ocean State

1. Problem Statement

Standard Transformers suffer from a structural limitation: self-attention incurs an $O(N^2)$ computational cost and requires an $O(N)$ memory cache (Key-Value cache) that grows linearly with context length. While recent works (RetNet, RWKV, Mamba, Titans) attempt to mitigate this by introducing new machinery like linear attention, selective gating, or external memory modules, they fundamentally alter the architecture. Furthermore, standard Transformers discard the final hidden state after every forward pass, despite it containing an absorption of the entire context.

The paper posits that the $O(N^2)$ cost is structural to the current implementation of attention, not inherent to the concept of context modeling. It asks: What if the hidden state, which absorbs the entire context, were not discarded but retained and used as the sole object of attention for the next token?

2. Methodology: The Ocean State

The proposed solution, "The Ocean State," involves two specific repositionings of the standard Transformer architecture without adding new components (e.g., no gating, no external memory).

2.1 Core Mechanisms

The Ocean State (Repositioning 1):
- Instead of discarding the hidden state after a forward pass, the output of the current token's processing replaces the previous hidden state, which is termed the "ocean state."
- Attention Mechanism: Attention is computed strictly between the current token and the persistent ocean state (a single dense vector).
- Consequences: This eliminates the need for Key-Value (KV) caches, positional encodings, and causal masks. The state is updated via a "relay" (1:1 replacement) rather than accumulation, ensuring information is rearranged but not compressed or discarded.
- Read-Only: The ocean state is read by the current token but never processed through the Feed-Forward Network (FFN) or layer stack; only the token passes through the stack.
Ocean BPTT (Repositioning 2):
- Because the ocean state is a recurrent variable, gradients must be unrolled through time (Backpropagation Through Time, BPTT).
- Gradient Flow: During the backward pass, the gradient of the ocean state at step $t$ is computed based on the token at step $t$ . This gradient is then added to the hidden state gradient of step $t-1$ .
- Read-Path Operator: A specific operator maps the token's gradient back to the ocean state using the attention weights saved from the forward pass. The more a token "reads" the ocean, the stronger the gradient signal sent back to update the ocean state.

2.2 Algorithmic Implementation

The training loop consists of:

Forward Pass: Token $t$ attends to $ocean_{t-1}$ . The token passes through the layer stack (Attention + FFN). The output replaces $ocean_{t-1}$ to become $ocean_t$ .
Backward Pass: Gradients flow from the loss at step $t$ back through the layers. Crucially, the gradient of the ocean state is accumulated across time steps, allowing errors from future tokens to update the representation of past contexts.

2.3 Engineering Variants

The paper explores three training regimes to balance cost, time, and quality:

Serial Ocean BPTT: Full time-chain unrolling. Highest quality gain, lowest compute cost, but serial wall-clock time.
Ocean on Framework: Runs the full chain on standard autograd frameworks. Maintains full quality gain but improves throughput via batching.
Ocean Detach: Cuts the time gradient chain (no BPTT across steps). Allows for full parallel training (like standard Transformers) but sacrifices the performance gain derived from time gradients.

3. Key Contributions

Architectural Repositioning: Demonstrates that $O(1)$ attention and $O(1)$ memory are achievable by reusing the discarded hidden state as a persistent "ocean" vector, eliminating the need for KV caches and positional encodings.
Emergent Capability: Shows that a Transformer can learn to read and extract information from a single dense vector without explicit instruction, relying purely on end-to-end training.
Scalability to Extreme Contexts: Proves the architecture can handle context lengths up to 10 million steps without degradation, a feat impossible for standard Transformers due to memory and compute constraints.
Gradient Unrolling: Introduces "Ocean BPTT," applying standard BPTT to the full layer stack output as the recurrent state, identifying this as the primary source of performance gains.

4. Experimental Results

Experiments were conducted on the Shakespeare corpus (character and word levels) using a custom NumPy implementation (float64) and verified with PyTorch.

4.1 Performance and Training Dynamics

Superiority: Across all tested configurations (widths $d \in \{32, 64, 128, 256\}$ , depths $L \in \{2, 4, 6\}$ , and various head counts), the "Ocean Full" regime consistently outperformed the standard Transformer.
Stable Gain: The performance gap ( $\Delta = Loss_{std} - Loss_{full}$ ) was positive and stable (approx. +0.15 loss reduction) across the grid.
Source of Gain: The gain is primarily structural. When the time chain is detached (Ocean Detach), the performance drops significantly, indicating the time-gradient unrolling is the critical factor, not just the attention mechanism itself.

4.2 Long-Context and Memory

No Forgetting: In a 10-million-step prefill test, the loss remained flat (1.532 at step 1M vs. 1.531 at step 10M). The ocean state showed no sign of capacity saturation or forgetting.
Echo Task: The model achieved near-zero loss (0.0006) in an echo task with a 2-million-step delay, whereas the standard Transformer failed completely.
Memory Efficiency: The ocean state requires only $O(1)$ memory (a single vector of dimension $d$ ), independent of context length. In contrast, standard Transformers require $O(N)$ memory for KV caches, which becomes infeasible at scales like 1 million tokens (requiring ~1 TB for a 13B model).

4.3 Inference and Generation

No Extrapolation: The ocean generates coherent text without needing positional encoding extrapolation (e.g., RoPE, ALiBi) because the context length is not an input variable to the attention mechanism.
Wall-Clock Time: While the serial nature of the recurrence makes training slower per step than parallel standard Transformers, the paper notes that layer-wise pipelining (Ocean A-2) can achieve a 1.9x speedup in wall-clock time, and nested parallelism (Ocean A-3) offers theoretical gains of $10^3$ . The author emphasizes that training involves an explicit trade-off triangle between cost, time, and quality, where practical implementation requires selecting specific regimes (e.g., detach for industrial scale, full chain for research quality).

5. Significance and Claims

The paper claims three primary points, which it defines as falsifiable:

Training Consistency: The Ocean state trains with the same learning curve shape as the standard Transformer, confirming it is a valid repositioning rather than a fundamentally different learning paradigm.
Universal Gain: For any Transformer architecture under identical conditions, the Ocean state yields a positive performance gain ( $\Delta > 0$ ).
Infinite Capacity: With sufficient dimension ( $d$ ), the Ocean state exhibits no capacity limit, no forgetting, and no retrieval failure at any observable scale.

The author concludes that the "Ocean State" addresses the principle of $O(1)$ inference and training for arbitrary context lengths, stating, "we believe the principle is solved." By relying on the emergent ability of the Transformer to organize information within a single vector, the method removes the structural bottlenecks of quadratic attention and linear memory growth, potentially enabling "small and personal" transformers and human-like robots with continuous cognitive states. The work emphasizes that these results are derived from "what-if" experimentation on existing components, respecting the emergent properties of the original architecture, while acknowledging open questions regarding specific configurations (e.g., why detach beats the standard in some cases).