File size: 4,478 Bytes
066b2f5 3f8f837 066b2f5 3f8f837 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | ---
license: mit
library_name: tensorflow
tags:
- othello
- reversi
- game-ai
- transformer
- mixture-of-experts
- rmsnorm
- swiglu
---
## 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](https://github.com/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](https://huggingface.co/rsu/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](https://huggingface.co/rsu/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
```python
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())
```
|