rsu/Reversi-Transformer-2

Transformer Reversi AI

This Reversi AI is a Transformer-based model that uses Mixture-of-Experts architecture for decision-making instead of CNNs.
Github - rsu-Suba/ReversiGPT

Architecture

Normal model M2.keras

• Parameters: 1.98M
• Model size: 40.4MB
• Input: (8, 8, 3)
  • Own stones
  • Opponent stones
  • Legal moves
• Board representation: 64 tokens
• Embedding dimension: 128
• Transformer blocks: 4
• Attention heads: 4
• MHA experts: 3 (with optional 8-direction attention mask)
• FFN experts: 2 (SwiGLU)
• Routing: Independent Top-k sparse routing with load balancing loss
• Iterative routing steps: 2
• Normalization: Pre-LN RMSNorm
• FFN activation: SwiGLU

Small model M2_small.keras

• Parameters: 756K
• Model size: 15.7MB
• Embedding dimension: 96
• Transformer blocks: 3
• Attention heads: 3
• MHA experts: 2
• FFN experts: 2
• Iterative routing steps: 1

Diagram

Input (8×8×3)
    ↓
64 Tokens
    ↓
Dense 128 (or 96)
    ↓
Token + Row + Column + Move Embeddings
    ↓
DynamicAssembly × 4 (or 3)
  ├─ MHA Expert × 3 (or 2) [MHA Router (Top-1)]
  └─ FFN Expert × 2 (SwiGLU) [FFN Router (Top-1)]
    ↓
    ├─ Policy Head → RMSNorm → Dense → 64 moves
    └─ Value Head  → Attention Pooling → Shared Dense
                        ├─ Win Head   → Tanh (Win Rate [-1, 1])
                        └─ Score Head → Tanh (Score Diff [-1, 1])

DynamicAssembly

Each DynamicAssembly block contains separate pools of MHA and FFN experts with independent routers. At each iteration step, the mean-pooled board representation is combined with a step embedding and passed to the respective router.

The router selects the top expert (Top-k sparse routing) and dispatches the computation, with auxiliary load-balancing loss during training to ensure uniform expert utilization.

Policy Head

The policy head applies RMSNorm and dense layers to predict a probability distribution over the 64 board positions.

Value Head

The value head uses Attention Pooling over the token sequence and predicts a 2-dimensional vector in [-1, 1] × [-1, 1]:

• Win Rate: Estimated game outcome [-1, 1]
• Score Difference: Estimated final stone difference [-1, 1]

Training

This model was trained using self-play game records generated by Reversi-Transformer-1.

Self-Play data generation

• MCTS simulations: 256
• MCTS batch size: 8
• PUCT constant: 2.5
• Total games: 30,000

Model Training

• Optimizer: AdamW
• Clipnorm: 1.0
• Weight decay: 0.01
• Learning rate: 2.0e-4
• Learning rate scheduler: WarmupCosineDecay
• Learning rate warmup: 2 epochs
• Training batch size: 512
• Epochs: 30

Evaluation

Reversi-Transformer-1 was used as the previous model.

Win rate

• vs Random : 100% (100 / 100 games)
• vs Prev model : 54% (54 / 100 games)

Average final stone count

• vs Random : 59
• vs Prev model : 38

Usage / Quick start

import numpy as np
import tensorflow as tf
from huggingface_hub import hf_hub_download
from model import RMSNorm, TokenAndPositionEmbedding, MHA, FFN, DynamicAssembly, AttentionPooling

model_path = hf_hub_download(repo_id="rsu/Reversi-Transformer-2", filename="M2.keras", repo_type="model")
# If you want to use M2_small
# model_path = hf_hub_download(repo_id="rsu/Reversi-Transformer-2", filename="M2_small.keras", repo_type="model")

custom_objects = {
    "RMSNorm": RMSNorm,
    "TokenAndPositionEmbedding": TokenAndPositionEmbedding,
    "MHA": MHA,
    "FFN": FFN,
    "DynamicAssembly": DynamicAssembly,
    "AttentionPooling": AttentionPooling,
}
model = tf.keras.models.load_model(model_path, custom_objects=custom_objects, compile=False)

board = np.zeros((1, 8, 8, 3), dtype=np.float32)
board[0, 3, 3, 1] = 1.0
board[0, 3, 4, 0] = 1.0
board[0, 4, 3, 0] = 1.0
board[0, 4, 4, 1] = 1.0

board[0, 2, 3, 2] = 1.0
board[0, 3, 2, 2] = 1.0
board[0, 4, 5, 2] = 1.0
board[0, 5, 4, 2] = 1.0

policy, value = model(board, training=False)
print("Policy (64 moves prob): ", policy.numpy())
print("Win rate est [-1, 1]: ", value[0, 0].numpy().item())
print("Score diff est [-1, 1]: ", value[0, 1].numpy().item())