TRM-Chess

Try it live: https://huggingface.co/spaces/octaviomartinez/trm-chess

English

A chess engine built with an experimental architecture: TRM (Tiny Recursive Model) combined with BitNet-style ternary weights. As far as could be confirmed, this combination had never been publicly tried before in the chess domain -- TRM and its predecessor HRM (Hierarchical Reasoning Model) were originally designed for grid-based tasks like Sudoku, mazes, and ARC-AGI, not classic board games.

Architecture

The core: TRM (Tiny Recursive Model). Unlike a network that sees the board once and responds, this model iterates over itself several times before deciding on a move: a recursive reasoning process with two internal states.

z: the reasoning "scratchpad", updated several times per cycle
y: the answer being refined, updated once per cycle

Both states are updated using the same shared network (2 transformer layers), not two separate networks -- that's the key simplification TRM introduced over HRM, which used two distinct networks for the same two functions.

Ternary weights (BitNet b1.58). The recursive core's layers quantize their weights to only three possible values: {-1, 0, +1}, with a separate scale factor (BitNet b1.58 style). This greatly reduces the core's computational cost. The output heads (policy and value) stay at normal precision, since they need to discriminate fine evaluation differences that quantization would ruin.

Input/output representation. Board: 19 planes of 8x8 (own/opponent pieces, turn, castling rights, en passant, halfmove clock), always from the mover's perspective. Moves: 64 origin squares x 64 destination squares x 4 promotion variants = 16,384 possible moves encoded as a single index.

Parameters: ~22 million total. The recursive core itself is small (consistent with TRM's spirit); most parameters live in the policy head, whose output size (16,384 classes) is inherently large regardless of core size.

Training

Method: Stockfish distillation. The model learned by imitating millions of real positions from the Lichess database, each with the move Stockfish considered best and its numeric evaluation -- not by playing against itself (self-play).

Dataset: Lichess/chess-position-evaluations (394M total positions available; a 12M subset filtered by Stockfish analysis depth >= 18 was used)
Optimizer: AdamW, with adaptive learning rate (drops only if the model stops improving) and automatic early stopping
Deep supervision: each reasoning cycle is trained with its own error signal, but the gradient only propagates through the last cycle (one-step approximation), avoiding memory proportional to all cycles combined
Label smoothing on the policy loss (softens the target when several moves are nearly equally good, a much more frequent issue in chess than in domains with a single correct answer like Sudoku)
Simple curriculum learning: early epochs see positions with fewer pieces before moving to more complex positions

Trained by combining free GPUs from Google Colab, Kaggle, and Lightning AI -- no dedicated or paid infrastructure.

How to use this model

!pip install huggingface_hub torch python-chess safetensors -q
!wget -q https://huggingface.co/octaviomartinez/trm-chess/raw/main/model_chess.py

import chess
import torch
from model_chess import TRMChess, board_to_tensor, get_legal_move_mask, index_to_move

model = TRMChess.from_pretrained("octaviomartinez/trm-chess")
model.eval()

board = chess.Board()
with torch.no_grad():
    out = model(board_to_tensor(board).unsqueeze(0))
    logits = out["policy_logits"].squeeze(0)
    mask = get_legal_move_mask(board)
    logits[~mask] = -1e9
    best_move = index_to_move(logits.argmax().item(), board)
print(best_move.uci())

For real play with MCTS search (stronger than raw policy alone), see play_chess.py in this repository.

Note on the "Use this model" widget: this repo doesn't use the standard transformers library integration (it uses the lighter huggingface_hub PyTorchModelHubMixin instead), so the automatic Transformers/Colab/Kaggle tabs shown on some other model pages don't appear here. The code snippet above works the same way -- just copy-paste it into any Python environment (Colab, Kaggle, or local).

MCTS search

For actual play, the model doesn't just use its raw policy -- it's combined with Monte Carlo Tree Search (MCTS), similar in spirit to AlphaZero: candidate moves are simulated, promising lines are explored deeper guided by the network's policy, and resulting positions are evaluated with the value head.

Measured playing strength

Measured against Stockfish limited to various Elo levels (using UCI_LimitStrength + UCI_Elo), playing real game series (not just self-play estimation). The real result landed considerably lower than initial self-play-based estimates, roughly 800-1200 Elo, evidencing a known limitation of pure distillation: the model generalizes worse against playing styles it never saw during training.

Honest limitations

Does not compete with professional-level engines (Stockfish, Leela Chess Zero). Those engines have years of development, tree-search-specialized architectures, and orders of magnitude more training data and compute.
No self-play phase -- all learning comes from imitating Stockfish, not from discovering its own strategies by playing itself. This likely explains much of the gap between self-play performance and performance against a genuinely different opponent.
May show indecision (repeating moves without converging) in positions converting material advantage into victory, particularly in long games.

Possible next steps

Self-play: training on games generated by the model itself (with real outcome as signal), instead of only imitating Stockfish.
More data and/or a larger recursive core.
MCTS guiding training itself, not just final play.

License

Apache 2.0.

Español

Motor de ajedrez construido con una arquitectura experimental: TRM (Tiny Recursive Model) combinado con pesos ternarios estilo BitNet. Hasta donde se pudo confirmar, esta combinacion nunca se probo publicamente antes en el dominio del ajedrez -- TRM y su predecesor HRM (Hierarchical Reasoning Model) se disenaron originalmente para tareas de grilla como Sudoku, laberintos y ARC-AGI, no para juegos de mesa clasicos.

Arquitectura

El nucleo: TRM (Tiny Recursive Model). A diferencia de una red que ve el tablero una vez y responde, este modelo itera sobre si mismo varias veces antes de decidir una jugada: un proceso de razonamiento recursivo con dos estados internos.

z: el "borrador" de razonamiento, se actualiza varias veces por ciclo
y: la respuesta que se va refinando, se actualiza una vez por ciclo

Ambos estados se actualizan usando la misma red compartida (2 capas transformer), no dos redes separadas -- esa es la simplificacion clave que propuso TRM sobre HRM, que usaba dos redes distintas para las mismas dos funciones.

Pesos ternarios (BitNet b1.58). Las capas del nucleo recursivo cuantizan sus pesos a solo tres valores posibles: {-1, 0, +1}, con un factor de escala aparte (estilo BitNet b1.58). Esto reduce mucho el costo computacional del nucleo. Las cabezas de salida (policy y value) se mantienen en precision normal, porque necesitan discriminar diferencias finas de evaluacion que la cuantizacion arruinaria.

Representacion de entrada y salida. Tablero: 19 planos de 8x8 (piezas propias/rivales, turno, derechos de enroque, captura al paso, contador de jugadas sin captura), siempre desde la perspectiva del jugador que mueve. Jugadas: 64 casillas de origen x 64 de destino x 4 variantes de promocion = 16,384 jugadas posibles codificadas como un unico indice.

Parametros: ~22 millones de parametros totales. El nucleo recursivo en si es chico (consistente con el espiritu de TRM); la mayor parte de los parametros vive en la cabeza de policy, cuyo tamano de salida (16,384 clases) es inherentemente grande sin importar el tamano del nucleo.

Entrenamiento

Metodo: destilacion de Stockfish. El modelo aprendio imitando millones de posiciones reales de la base de datos de Lichess, cada una con la jugada que Stockfish considero mejor y su evaluacion numerica -- no jugando contra si mismo (self-play).

Dataset: Lichess/chess-position-evaluations (394M de posiciones totales disponibles; se uso un subconjunto de 12M filtrado por profundidad de analisis de Stockfish >= 18)
Optimizador: AdamW, con learning rate adaptativo (baja solo si el modelo deja de mejorar) y early stopping automatico
Deep supervision: cada ciclo de razonamiento se entrena con su propia senal de error, pero el gradiente solo se propaga a traves del ultimo ciclo (aproximacion de un paso), para no necesitar memoria proporcional a todos los ciclos juntos
Label smoothing en la perdida de policy (suaviza el objetivo cuando hay varias jugadas casi igual de buenas, un problema mucho mas frecuente en ajedrez que en dominios con una unica respuesta correcta como Sudoku)
Curriculum learning simple: las primeras epocas ven posiciones con menos piezas antes de pasar a posiciones mas complejas

Se entreno combinando GPUs gratuitas de Google Colab, Kaggle, y Lightning AI -- sin infraestructura dedicada ni paga.

Como usar este modelo

!pip install huggingface_hub torch python-chess safetensors -q
!wget -q https://huggingface.co/octaviomartinez/trm-chess/raw/main/model_chess.py

import chess
import torch
from model_chess import TRMChess, board_to_tensor, get_legal_move_mask, index_to_move

model = TRMChess.from_pretrained("octaviomartinez/trm-chess")
model.eval()

board = chess.Board()
with torch.no_grad():
    out = model(board_to_tensor(board).unsqueeze(0))
    logits = out["policy_logits"].squeeze(0)
    mask = get_legal_move_mask(board)
    logits[~mask] = -1e9
    best_move = index_to_move(logits.argmax().item(), board)
print(best_move.uci())

Para jugar de verdad con busqueda MCTS (mas fuerte que solo la policy cruda), ver play_chess.py en este mismo repositorio.

Nota sobre el panel "Use this model": este repo no usa la integracion estandar de la libreria transformers (usa en cambio el mixin mas liviano PyTorchModelHubMixin de huggingface_hub), asi que las pestanas automaticas de Transformers/Colab/Kaggle que se ven en otras paginas de modelos no aparecen aca. El codigo de arriba funciona igual -- solo hay que copiarlo y pegarlo en cualquier entorno de Python (Colab, Kaggle, o local).

Busqueda con MCTS

Para jugar, el modelo no usa solo su policy directamente -- se combina con una busqueda tipo Monte Carlo Tree Search (MCTS), similar en espiritu a AlphaZero: se simulan muchas jugadas candidatas, se profundiza en las lineas mas prometedoras guiado por la policy de la red, y se evaluan las posiciones resultantes con el value head.

Fuerza de juego medida

Se midio contra Stockfish limitado a distintos niveles de Elo (usando UCI_LimitStrength + UCI_Elo), jugando series de partidas reales (no solo estimando por partidas contra si mismo). El resultado real quedo en un rango bastante mas bajo que las estimaciones iniciales basadas en partidas contra si mismo, aproximadamente 800-1200 Elo, evidenciando una limitacion conocida de la destilacion pura: el modelo generaliza peor contra estilos de juego que nunca vio durante el entrenamiento.

Limitaciones honestas

No compite con motores de nivel profesional (Stockfish, Leela Chess Zero). Esos motores tienen anos de desarrollo, arquitecturas especializadas en busqueda de arbol, y entrenamiento con ordenes de magnitud mas de datos y computo.
No hubo fase de self-play -- todo el aprendizaje viene de imitar a Stockfish, no de descubrir estrategias propias jugando contra si mismo. Esto probablemente explica buena parte de la brecha entre el desempeno contra si mismo y contra un rival genuinamente distinto.
Puede mostrar indecision (repetir jugadas sin converger) en posiciones de conversion de ventaja material a victoria, particularmente en partidas largas.

Proximos pasos posibles

Self-play: entrenar sobre partidas generadas por el propio modelo (con resultado real como senal), en vez de solo imitar a Stockfish.
Mas datos y/o un nucleo recursivo mas grande.
MCTS guiando el entrenamiento mismo, no solo el juego final.

Licencia