Spatial Wave Transformer (SWT)

spatialwave is a Python library that implements an experimental Artificial Intelligence architecture based on a 3D spatial field. Instead of global attention mechanisms, information propagates through local 3D updates, energy-gated resonance, and dynamic causal wormhole shortcuts.

The library is designed for research, experimentation, and proof-of-concept AI development. It natively supports GPU/CUDA acceleration and works on both PyTorch CPU/GPU runtimes.

Key Features

• 3D Latent Field representation: Model state resides in a tridimensional volume where information maps to specific coordinates.
• Causal 3D Convolutions: Multi-cycle propagation preserves sequence order dynamically via raster-scan causality masks.
• Resonance Gating: Zones of high energy gate the computation, driving focus onto highly active signal junctions.
• Causal Wormhole Layers: Dynamic, sparse shortcuts connect highly active regions using causal Multi-Head Attention.
• Iterative Reasoning Engine: Executes configurable cycles with adaptive stopping, shutting down iterations when states converge.
• Word-Based BasicTokenizer: Full fitting, encoding, decoding, and serialization out-of-the-box.
• Robust Trainer: Simplified training loop with progress tracking, perplexity computation, and automatic checkpoints.

Installation

To install the library, clone this repository and run:

git clone https://github.com/bueormnew/spatialwave.git
cd spatialwave
pip install -e .

Or install dependencies manually:

pip install -r requirements.txt

Quickstart Example

Here is a simple example showing how to initialize, train, save, load, and run inference using spatialwave:

from spatialwave import SpatialWaveModel, BasicTokenizer, TextGenerator

# 1. Prepare some text
corpus = [
    "El espacio tridimensional contiene informacion latente.",
    "La resonancia local ayuda a propagar las ondas de datos.",
    "Los agujeros de gusano conectan regiones distantes causales."
]

# 2. Tokenize text
tokenizer = BasicTokenizer()
tokenizer.fit(corpus)

# 3. Create Model
model = SpatialWaveModel(
    vocab_size=tokenizer.vocab_size,
    field_size=16,       # 16x16x16 volume
    latent_dim=64,      # Hidden size
    cycles=6,           # Max cycles
    wormhole_interval=2 # Wormhole every 2 cycles
)

# 4. Generate continuation
generator = TextGenerator(model, tokenizer, device="cpu")
print(generator.generate("El espacio", max_tokens=5))

Training a Model

You can use the built-in Trainer to run autoregressive training on a text corpus:

from spatialwave import Trainer, TextDataset

# Create dataset
dataset = TextDataset(corpus, tokenizer, max_seq_len=16)

# Create trainer
trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    device="cpu",
    checkpoint_dir="checkpoints"
)

# Train model
trainer.train(epochs=10, batch_size=4, lr=1e-3)

Save and Load

You can save and load both the model and tokenizer to easily resume work or run inference elsewhere:

# Save model & tokenizer
model.save("model.swt")
tokenizer.save("tokenizer.json")

# Load model & tokenizer
loaded_model = SpatialWaveModel.load("model.swt")
loaded_tokenizer = BasicTokenizer.from_file("tokenizer.json")

Internal Architecture

The SWT architecture processes tokens as follows:

1. Input Injection: Text tokens are embedded and injected into the volume along a raster-scan grid.
2. Local Propagation: 3D causal convolutions distribute information to adjacent neighbors ($3 \times 3 \times 3$ grid).
3. Resonance Gating: Energy updates dynamically. High-energy zones gate state updates.
4. Wormholes: Multi-Head Attention runs on top-K active cells (sorted causally) to exchange distant state.
5. Iterative Loops: The cycles run repeatedly until convergence or max cycle threshold.
6. Logits Prediction: Logits are extracted from the token coordinates and mapped to vocabulary predictions.

Read the detailed Architecture Documentation for more information.

Roadmap

• Advanced Tokenization: Support Byte-Pair Encoding (BPE) and SentencePiece.
• Multimodality: Support image embedding injection on boundary faces (e.g. front face) and text on another.
• Longer Sequences: Optimized sparse 3D convolutions for scaling grid size.
• Multi-GPU Training: Native PyTorch DistributedDataParallel (DDP) support.