kbeta-transformer2d – 2-D Heat-Diffusion Transformer trained with Kourkoutas-β 🌞🦎🚀📈

Research companion code for the upcoming paper
“Kourkoutas-β: A Sunspike-Driven Adam Optimizer with Desert Flair.”
Published as arXiv:2508.12996.

This repository contains the full 2‑D data‑driven Transformer workload that accompanies the optimiser
(see the separate kbeta repo), plus lightweight utilities for training, evaluation and visualisation.

Table of Contents

1. Why a 2-D Transformer?
2. Model highlights
3. Project layout
4. Installation
5. Quick start
6. Command-line interface
7. Training from scratch
8. Using your own datasets
9. Tests & linting
10. Relation to Kourkoutas-β
11. Learning-rate schedule behaviour
12. Citation
13. License

Why a 2‑D Transformer?

• Spatial‑temporal diffusion appears in countless engineering problems (heat flow, pollutant transport, …).
• A purely data‑driven Transformer offers a clean stress‑test for the optimiser.
• Solver‑free physics loss: we embed the heat‑equation residual as an analytic term, no back‑prop through external PDE solvers is required.
• The model scales to 512 × 512 meshes on Apple Silicon while remaining <2 M parameters; perfect for rapid experimentation.

Model highlights

(what’s special about HeatDiffusion-Transformer-2D)

• Patch-wise attention on 2-D grids
  The input tensor is reshaped into (T × H × W) patches, letting the model treat every spatial location symmetrically while still exploiting MX-GPU tensor cores efficiently.

• Dual masking modes
  Causal masks give an autoregressive model useful for long-horizon rollout tests; block masks allow full-context training when future frames are available.

• RoPE (Rotary Positional Encoding) in the time dimension
  A single line swap lets you switch between vanilla sinusoidal encodings and RoPE, which markedly improves extrapolation beyond the training window.

• Activation quantisation ready
  All dense / conv projections are implemented with mlx.nn.quantize_lin, giving you 8-bit weights on Apple Silicon without code changes.

• Paper configuration ≈ 32 M parameters
  With 24 encoder layers, 16 heads, embed_dim=512, and mlp_dim=256, the model has about 32 M trainable parameters—large enough to stress-test optimisers, yet compact enough to train on a single Mac Studio GPU (batch size 4, grid 25 × 25 × 401).

• One‑liner optimiser swap
  The model inherits its optimiser object, so comparing Adam vs Kourkoutas‑β is literally one YAML entry.

Project layout

kbeta-transformer2d
├── src/kbeta_transformer2d/
│   ├── data.py              # mesh generation + loaders
│   ├── model.py             # Transformer & loss
│   ├── optim_factory.py     # Kourkoutas‑β wiring
│   ├── train.py             # training / eval loops
│   ├── plot_utils.py        # visualisations
│   └── demo_heat2d.py       # CLI entry‑point
├── configs/
│   └── heat2d.yml           # default hyper‑params
│   └── paper.yml            # paper   hyper‑params
└── README.md                # you are here

Installation

Option 1: PyPI wheels (end‑users)

If you only want to run the Transformer benchmark with the latest kbeta:

pip install kbeta-transformer2d

For dev tools and tests:

pip install "kbeta-transformer2d[dev]"

For exact reproducibility of the paper (MLX 0.26.3 + pinned deps):

pip install "kbeta-transformer2d[repro]"

Option 2: Cloning the repo (researchers / contributors)

git clone https://github.com/sck-at-ucy/kbeta-transformer2d.git
cd kbeta-transformer2d
python -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"

This makes all configs and scripts editable for research use.

Quick start

pytest -q   # ➜ all smoke‑tests should pass

Command‑line interface

demo_heat2d.py is both import‑able and executable:

python -m kbeta_transformer2d.demo_heat2d <CONFIG.yml> [flags]

Notes on spike tracking
To actually record Sun‑spike/β₂ you need all of: --optimizer=kourkoutas, --kour_diagnostics, and --collect_spikes. Enabling --collect_spikes auto-enables --kour_diagnostics as well. The windowing/plot stride is controlled via YAML (see below).

Examples

# train 5 epochs with vanilla Adam‑(0.9,0.95)
python -m kbeta_transformer2d.demo_heat2d heat2d.yml \
   --epochs=5 \
   --optimizer=adam95 \
   --override storage.outdir="./OUTPUTS/run1"

# same as above but change mesh size and disable plotting
python -m kbeta_transformer2d.demo_heat2d heat2d.yml \
  --override geometry.dx=0.08 geometry.dy=0.08 \
  --override viz.enabled=false storage.outdir="./OUTPUTS/run2"  

# run with the *packaged* default config (no file in cwd needed)
python - <<'PY'
import subprocess, importlib.resources as res
cfg = res.files("kbeta_transformer2d.configs") / "heat2d.yml"
subprocess.run([
    "python", "-m", "kbeta_transformer2d.demo_heat2d",
    str(cfg),
    "--epochs=1",
    "--override", "storage.outdir=./OUTPUTS/run3"
])
PY

Tip: Use --override storage.outdir=... to redirect checkpoints/plots to a dedicated folder instead of cluttering the repo root.

Example configuration (excerpt)

seed: 0
geometry:
  rod_length: 1.0
  rod_width: 1.0
  dx: 0.04
  dy: 0.04
boundary_conditions:
  left_limits: [0, 1]
  right_limits: [0, 1]
  top_limits: [0, 1]
  bottom_limits: [0, 0.1]
thermal_diffusivity:
  alpha_limits: [0.01, 0.1]
model_params:
  start_predicting_from: 5
  batch_size: 4
  epochs: 10
  time_steps: 401
  num_heads: 16
  num_encoder_layers: 24
  mlp_dim: 256
  embed_dim: 512
  mask_type: block
learning_rate_schedule:
  5: 1.0e-3
  30: 5.0e-4
  40: 1.0e-4
  60: 1.0e-5
 120: 1.0e-6 
optimizer:
  name: adam999
  init_lr: 1.0e-3
  target_lr: 1.0e-5
  ramp_steps: 60000
tracking:
  collect_spikes: false     # set true to gather Sun‑spike / β₂ (if using kbeta)
  window: 500               # epochs per accumulation window
  plot_stride: 5000         # violin sampling stride (defaults to 10×window)
storage:
  outdir: null              # default = CWD; creates structured sub‑folders
io_and_plots:
  model_saving: false       # save a *final* full checkpoint at the end
save_checkpoints: true      # periodic checkpoints during training
save_interval: 10           # only used if save_checkpoints=true

(see configs/heat2d.yml and configs/paper.yml for the full list)

Learning‑rate schedule behaviour

• If a learning_rate_schedule block is present in your YAML config, the model will use that explicit step schedule (this is how the published paper runs were done).
• If no learning_rate_schedule is defined, the code will fall back to a cosine schedule controlled by init_lr, target_lr, and ramp_steps under the optimizer block.

YAML quick-reference — common pitfalls 🔍

betas:\n  - 0.9\n  - 0.999

Tip: If you’re ever unsure how YAML will parse a value, run
python -c 'import yaml, sys, pprint, pathlib; pprint.pprint(yaml.safe_load(pathlib.Path("your.yml").read_text()))'
to see exactly what Python receives.

Training from scratch

python -m kbeta_transformer2d.demo_heat2d configs/heat2d.yml --epochs=30

Checkpoints (.pkl + .safetensors) and plots are written under a structured OUTPUTS folder. A typical layout looks like:

$\text{text} \text{OUTPUTS}/ └── <\text{run\_label}>\text{\_}<\text{strategy}>\text{\_}<\text{mask}>/ ├── \text{datasets}/ # \text{saved} \text{MLX}/\text{NumPy} \text{arrays} ├── \text{checkpoints}/ # \text{periodic} \text{checkpoints} (\text{if} \text{enabled}) ├── \text{frames}/ # \text{prediction} \text{frames} / \text{movies} (\text{optional}) ├── \text{sunspike\_violin}/ # \text{violin} + \text{swarm} \text{plots} ├── \text{beta2\_violin}/ ├── \text{sunspike\_density}/ # \text{value} \times \text{epoch} \text{heat}‑\text{maps} ├── \text{beta2\_density}/ └── \text{mse}/ # \text{block}/\text{AR} \text{MSE} \text{curves} $

Checkpoint policy:

• If save_checkpoints: false, no periodic checkpoints are written (we internally treat save_interval=None).
• If true, we save every save_interval epochs (default = 10).

Using your own datasets

1. Provide your dataset as NumPy/MLX arrays.
2. Adjust geometry.* in the YAML to match mesh resolution.
3. Replace or extend generate_datasets() in data.py if needed.

Tests & linting

pytest
ruff check .
mypy src

Relation to Kourkoutas‑β

This repo uses the optimiser from kbeta; it does not re‑implement it.
optim_factory.py wires KourkoutasBeta into the training loop.

Further Reading & Related Resources 📚

Citation

If you use this work, please cite both the paper and the software archive:

Paper (arXiv preprint)

@article{Kassinos2025Kourkoutas,
  title   = {Kourkoutas-β: A Sunspike-Driven Adam Optimizer with Desert Flair},
  author  = {Stavros Kassinos},
  journal = {arXiv preprint arXiv:2508.12996},
  year    = {2025},
  url     = {https://arxiv.org/abs/2508.12996}
}

Software (Zenodo archive)

@software{kassinos2025transformer2d,
  author       = {Stavros Kassinos},
  title        = {kbeta-transformer2d: 2-D Heat-Diffusion Transformer – Companion Code},
  year         = {2025},
  publisher    = {Zenodo},
  version      = {1.0.0},
  doi          = {10.5281/zenodo.16911539},
  url          = {https://doi.org/10.5281/zenodo.16911539}
}

License

MIT. See LICENSE for the full text.

Happy experimenting — and may your gradients be sunny ☀️🦎🚀📈