NeuroSymbolicCodeTransformer

A 13.32M-parameter hybrid neuro-symbolic transformer that learns Python code generation from mathematical representations of Abstract Syntax Trees (ASTs). The model distills knowledge from a 9B-parameter reasoning LLM (Ornith/Qwen3.5-based) into a compact architecture trainable on consumer GPUs with only 6GB VRAM.

📄 Whitepaper: NeuroSymbolicCodeTransformer: A Neuro-Symbolic Transformer for Code Generation via Knowledge Distillation and REINFORCE Policy Gradient (PDF, 12 pages)

🌐 GitHub Pages: dexopt.github.io/neuro-symbolic-coder-13M

Overview

The NeuroSymbolicCodeTransformer uses a three-phase pipeline: (1) a teacher LLM converts natural language requests into mathematical specifications, (2) a small neuro-symbolic model learns to map those specifications to AST token sequences via supervised learning, and (3) REINFORCE policy-gradient reinforcement learning refines the model using teacher-scored feedback.

The architecture introduces five custom neuron classes (LogicGate, LocalTree, Program, Routing, Memory) organized into MoE-routed transformer blocks. The entire model occupies only 25.4 MiB of VRAM at bfloat16 precision.

Architecture

                                  NeuroSymbolicCodeTransformer Pipeline
                          ==============================================================

                              +------------+          +------------------+
  User Request (English) ---->|  Ornith 9B  |-------->|  Math Spec JSON  |
                              |  (Teacher)  |         +------------------+
                              +------------+                 |
                                     ^                       v
                                     |              +------------------+
                                     |              |  Token Encoder   |
                                     |              | (125+ vocab      |
                                     |              |  tokens)         |
                                     |              +------------------+
                                     |                       |
                                     |                       v
                                     |              +------------------+
                                     |              | NeuroSymbolic    |
                                     |              | Code Transformer |
                                     |              | (13.32M params)  |
                                     |              +------------------+
                                     |                       |
                                     |                       v
                                     |              +------------------+
                                     |              |  Token Decoder   |
                                     |              | (AST tokens)     |
                                     |              +------------------+
                                     |                       |
                                     |                       v
                                     |              +------------------+
                                     |              |  ast_to_code     |
                                     |              |  (Python render) |
                                     |              +------------------+
                                     |                       |
                              +------------+               v
                              | Validation |<---- Generated Python Code
                              |  (Score)   |
                              +------------+

                    Training Pipeline (3 Phases)
                    ===========================

  Phase 1: Data Generation    Phase 2: Supervised        Phase 3: REINFORCE RL
  -------------------------   -----------------------    --------------------------
  Prompt --> Ornith 9B        math_spec --> token enc      random spec generation
     |                               |                        |
     v                               v                        v
  Code --> MathEncoder         NeuroSymbolicTransformer  generate 2 candidates
     |                               |                        |
     v                               v                        v
  MathSpec --> TokenEncoder     cross-entropy loss       Ornith scores each
     |                                                     |
     v                                                     v
  Training Dataset                                   best > threshold?
                                                    yes --> policy gradient
                                                    no  --> skip

Transformer Block Detail

  Input Tensor (B, S, D=320)
        |
  +-----v------+
  | MultiHead  |  4 heads, head_dim=80, batch_first
  | Attention  |  dropout=0.1
  +-----+------+
        |
  +-----v------+
  | LayerNorm  |
  +-----+------+
        |
  +-----v------+
  |  Routing   |  Mean-pool --> 2-layer MLP --> 4-way softmax
  |  Neuron    |  noise injection during training
  +-----+------+
        |
  +-----v------------------------------------+
  |       Expert Bank (weighted sum)          |
  |                                           |
  |  [LogicGateNeuron]  [ProgramNeuron]       |
  |  [LocalTreeNeuron]  [Linear Fallback]     |
  |                                           |
  +-----+------------------------------------+
        |
  +-----v------+
  | LayerNorm  |
  +-----+------+
        |
  +-----v------+
  |  Memory    |  32-slot differentiable memory
  |  Neuron    |  attention read + gated write
  +-----+------+
        |
  Output (B, S, D=320)

Neuron Architecture

Neuron	Dimensions	Purpose	Mechanism
LogicGateNeuron	d_model -> d_model/2 -> d_model	Conditional AST nodes (if/else, boolean ops)	Projects input into two operand paths, applies AND/OR/XOR/NOT logic with learnable softmax weighting
ProgramNeuron	d_model -> d_model; 5 op chunks	Differentiable compute primitives	5 parallel math operations (sum, product, max, min, mean) per chunk; learnable softmax over operations
LocalTreeNeuron	d_model; kernel=3	Parent-child AST hierarchy	Depthwise Conv1d with sigmoid gating; local windows simulate tree branches
RoutingNeuron	d_model -> d_model/4 -> 4	Expert dispatch	MoE-style gating: mean-pooled representation through 2-layer MLP to 4-way softmax; noise added during training
MemoryNeuron	32 slots x d_model	Long-range dependency memory	Attention-based read over 32-slot memory bank; gated soft-update writes during training (decay=0.99, rate=0.01)

Model Configuration

Parameter	Value
Total parameters	13,320,345
d_model (hidden dimension)	320
Number of heads	4
Head dimension	80
Number of layers (blocks)	5
Vocab size	5,000
Max sequence length	512
Memory slots per block	32
Dropout	0.1
Expert count	4 (logic, program, tree, linear)
Precision	bfloat16 AMP
Gradient checkpointing	Yes (use_reentrant=False)

Training Results

Run A: Remote Ornith 9B (Successful)

Teacher Model:

Location: LM Studio at 192.168.1.102:1234
Model: ornith-1.0-9b-oq4 (Qwen3.5-based, 9B params, reasoning model)
Temperature: 0.6, top_p: 0.95

Data Generation:

Prompts: 6 of 8 examples generated successfully via chat/completions
Data generation time: ~8 minutes

Supervised Training:

Epochs: 5 (early stopping, patience=3)
Batch size: 4, gradient accumulation steps: 4
Loss: 8.63 -> 8.62
Training time: ~1 minute
Precision: bfloat16 AMP

Reinforcement Learning:

Algorithm: REINFORCE policy gradient
Iterations: 30
Candidates per iteration: 2
Score threshold: 70.0

RL Performance:

Metric	Value
Best score	100.00 / 100
Average score	28.33 / 100
RL training time	~59 minutes
Total training time	~68 minutes

VRAM Usage:

Metric	Value
Model memory (bf16)	25.4 MiB
Free VRAM	4.9 GiB / 6.0 GiB

Run B: Local Ornith (Partial)

Configuration:

Location: LM Studio at 127.0.0.1:1234
Prompts: 20 diverse prompts, cycling to 1,000 total examples
Example generation time: ~8-10 seconds each

Outcome:

First 12 examples succeeded (HTTP 200, ~8-10s per example)
Example 13 onwards: repeated LM Studio crashes
Error signature: "Engine protocol predict request failed: fetch failed"
Secondary error: "Model is unloaded"
Root cause: LM Studio server instability under sustained generation load

Run A (remote server) completed the full training pipeline. Run B (local server) failed during data generation and is documented here for reproducibility reference.

Pipeline Components

Component	File	Description
Configuration	`config.py`	Typed dataclass config with env presets (dev/training/inference), serialization
Math Encoder	`math_encoder.py`	Python source -> structured math spec via recursive AST traversal, 50+ node types
Token Encoder	`token_encoder.py`	MathSpec -> flat token sequences, 125+ structural markers, deterministic round-trip
Neuro-Symbolic Model	`small_model.py`	13.32M param transformer with 5 custom neuron classes, gradient checkpointing
Training Pipeline	`training_pipeline.py`	Supervised + REINFORCE RL with LM Studio as teacher, batch training
Inference Service	`inference_service.py`	FastAPI server with /generate, /validate, /batch, /health endpoints
Inference CLI	`infer.py`	Command-line inference script
Orchestrator	`main.py`	Training entry point, device detection, checkpoint management
Test Suite	`test_suite.py`	27 unit/integration tests
Memory Utils	`memory_utils.py`	VRAM monitoring, gradient accumulation helpers, AMP setup
Requirements	`requirements.txt`	PyTorch 2.0+, FastAPI, openai, uvicorn, pydantic, tqdm
Setup	`setup.sh`	One-command environment setup

Codebase stats: 12 files, ~4,300 lines of Python, 27 tests (100% pass).

Hardware Requirements

Resource	Minimum	Recommended
GPU	NVIDIA RTX 4050 (6GB VRAM)	NVIDIA RTX 4060+ (8GB+ VRAM)
VRAM peak	~500 MiB (bf16 AMP + gradient checkpointing)	~1 GiB (with safety margin)
System RAM	8 GB	16 GB
CPU	Any modern x86-64	AMD Ryzen / Intel Core i5+
Storage	2 GB	5 GB (checkpoints + logs)

Software requirements:

Python 3.10+
PyTorch 2.0+ with CUDA 12.1
LM Studio (for teacher model)

Quick Start

# Clone and setup
chmod +x setup.sh
./setup.sh

# Or manually
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

# Verify installation
python test_suite.py

LM Studio Setup

The teacher LLM runs separately in LM Studio. The pipeline communicates with it via the OpenAI-compatible API.

Download and install LM Studio from https://lmstudio.ai/
Load a reasoning-capable model. Recommended: ornith-1.0-9b-oq4 (Qwen3.5-based, 9B parameters)
Configure the model: temperature=0.6, top_p=0.95
Start the local inference server on port 1234
Verify connectivity:

curl http://localhost:1234/v1/models

For Run A, the teacher ran on a remote machine at 192.168.1.102:1234. Update config.py or pass --server to point to your LM Studio instance:

python main.py --server http://192.168.1.102:1234

Training

# Full training pipeline (default: localhost:1234)
python main.py --server http://<lm-studio-host>:1234

# With custom model name
python main.py --server http://localhost:1234 --model ornith-1.0-9b

Training proceeds through three phases:

Data generation (Phase 1): LM Studio generates Python code from prompts. CodeToMathConverter produces AST specs, MathSpecTokenizer encodes to token sequences.
Supervised learning (Phase 2): Model learns (math_spec -> math_spec) pairs with cross-entropy loss, bfloat16 AMP, gradient accumulation.
REINFORCE RL (Phase 3): Random spec generation, 2 candidate codes per iteration, teacher scores each. Parameter update when best candidate exceeds threshold.

Checkpoints are saved to ./checkpoints/.

Configuration

Edit config.py or create a JSON config file:

from config import Config

# Defaults (training env)
cfg = Config()
cfg.model.d_model = 320        # Transformer hidden dimension
cfg.model.n_layers = 5         # Number of blocks
cfg.model.n_heads = 4          # Attention heads
cfg.training.batch_size = 4    # Per-step samples
cfg.training.epochs = 20       # Supervised epochs
cfg.training.gradient_accumulation_steps = 4
cfg.rl.num_iterations = 20     # REINFORCE iterations
cfg.rl.num_candidates = 2      # Candidates per iteration
cfg.rl.score_threshold = 70.0  # Minimum score to trigger update

Inference

Start the FastAPI inference server:

python inference_service.py

The API runs on http://localhost:8000.

Endpoints

Method	Path	Description
`POST`	`/generate`	Generate code from natural language specification
`POST`	`/validate`	Validate existing code against a specification
`POST`	`/batch`	Batch code generation
`GET`	`/health`	Service and LM Studio availability check

Example Requests

Generate code from natural language:

curl -X POST http://localhost:8000/generate \
  -H "Content-Type: application/json" \
  -d '{
    "specification": "Create a function that filters even numbers from a list",
    "temperature": 0.7
  }'

Validate existing code:

curl -X POST http://localhost:8000/validate \
  -H "Content-Type: application/json" \
  -d '{
    "code": "def f(x): return [i for i in x if i % 2 == 0]",
    "specification": "filter even numbers from list"
  }'

Health check:

curl http://localhost:8000/health

CLI Inference

python infer.py --prompt "Write a function to compute fibonacci numbers" --checkpoint ./checkpoints/best.pt

How It Works

Neuro-Symbolic Approach

Instead of learning raw source code text, the model learns from mathematical representations of ASTs. The pipeline:

Encoding: Python code is parsed to an AST, then converted to a structured mathematical specification (nested dicts with typed operators, control flow markers, value annotations).
Tokenization: The math spec is flattened to a token sequence with structural boundary markers (LIST_START, DICT_START, KEY_SEP) enabling deterministic round-trip encoding.
Learning: The transformer block routes each input through an MoE bank of neuro-symbolic experts: LogicGateNeuron handles conditionals, LocalTreeNeuron captures parent-child AST structure, ProgramNeuron models differentiable computation, and a linear fallback handles unstructured patterns.
Memory: Each block ends with a 32-slot differentiable memory neuron that reads via attention and writes via gated soft-updates, providing long-range dependency tracking.
Decoding: Generated token sequences are deterministically decoded back to math specs, then rendered to Python source code.

Gradient Checkpointing

Each NeuroSymbolicTransformerBlock uses torch.utils.checkpoint.checkpoint with use_reentrant=False to trade compute for memory. This keeps peak VRAM at ~500 MiB despite the full model graph.

Reproducibility

Run A (Successful Full Training)

Requirements for reproduction:

LM Studio on a remote machine running ornith-1.0-9b-oq4 at port 1234
NVIDIA RTX 4050 (6GB) or better with CUDA 12.1
PyTorch 2.0+, Python 3.10+

python main.py --server http://<remote-lmstudio-ip>:1234 --model ornith-1.0-9b

Settings used: supervised epochs=5, patience=3, batch_size=4, grad_accum=4, bf16 AMP. RL iterations=30, candidates=2.

Run B (Data Generation Only, Unstable)

Run B attempted 1,000 examples from 20 cycled prompts on a local LM Studio instance. The server crashed repeatedly after 12 successful examples due to "Engine protocol predict request failed: fetch failed" errors. This appears to be an LM Studio stability issue under sustained chat/completions load. Using a remote server (Run A setup) resolves the problem.

Known Limitations

The 9B teacher model requires a separate machine or at least 8GB VRAM for inference
The 13.32M student model produces only ~500 MiB peak VRAM usage and is intentionally compact
Math spec vocabulary is Python-specific; the architecture generalizes to any structured-language AST
RL training time is dominated by teacher scoring latency (~59 min for 30 iterations)
LM Studio server stability varies: remote deployment recommended over local for sustained generation

Citation

@software{NeuroSymbolicCodeTransformer2026,
  author = {},
  title = {NeuroSymbolicCodeTransformer: A Neuro-Symbolic Transformer for Code Generation from AST Mathematical Specifications},
  year = {2026},
  url = {},
  note = {13.32M parameters, trained with knowledge distillation from Ornith 9B via REINFORCE RL on RTX 4050},
  keywords = {neuro-symbolic, code-generation, transformer, reinforcement-learning, mixture-of-experts, knowledge-distillation}
}

License

MIT License. See the repository for full terms.

Downloads last month: -; Downloads are not tracked for this model. How to track