LatentCalc: Rational Arithmetic via Deep Equilibrium and Universal Transformers

LatentCalc is a PyTorch-based, sequence-to-sequence model designed to solve rational arithmetic operations in latent space. The project features a hybrid "Sandwich" architecture, offering a selectable reasoning core: Universal Transformers with Adaptive Computation Time (ACT) or Deep Equilibrium Models (DEQ) using implicit differentiation.

Key Features

• Selectable Reasoning Core:
  • ACT (Adaptive Computation Time): Recurrent unrolling using Backpropagation Through Time (BPTT) with dynamic halting probabilities.
  • DEQ (Deep Equilibrium Models): Solves directly for the latent steady-state representation z^* = f(z^*, x) using implicit root-solving.
• Anderson Acceleration:
  • Features a custom, Tikhonov-regularized Anderson Acceleration solver using normal equations (compatible with Apple MPS/Mac, CPU, and CUDA/Nvidia GPUs) to dramatically accelerate fixed-point convergence.
• Stabilized DEQ Training:
  • Latent State Noise Injection: Optionally perturbs the hidden states with zero-mean Gaussian noise (--noise-sigma) during training to enforce contractivity and smooth the loss landscape.
  • Orthogonal Weight Initialization: Stabilizes early convergence by maintaining singular values at 1.
  • Core Spectral Normalization: Restricts the recurrent core's weight matrices to spectral norm \le 1 to guarantee contractive iterations.
  • Split Learning Rates: Configures separate parameter groups in AdamW (slower learning rate for the core reasoning block to prevent representation drift).
  • Weight Decay: Regularization to keep weights bounded.
• Accelerated Backpropagation:
  • Configurable Taylor/Neumann series approximation orders (--taylor-order 0, 1, etc.) to skip expensive adjoint fixed-point solving during the backward pass.
• Flexible Core Modules:
  • Selectable core blocks (--core-type transformer or mamba) to compare Attention-based models with pure-PyTorch Structured State Space (S4D) selective scans.
• High Performance on CUDA:
  • Automatic activation of TensorFloat-32 (TF32) precision for float32 matrix multiplications.
  • Support for torch.compile to fuse operations and accelerate execution.
• Curriculum Learning:
  • Generates fraction mathematical operations (+, *) dynamically on-the-fly, graduating the target coefficient range limit \mu (doubling it) and halving the solver tolerance threshold $\epsilon$ whenever the validation loss falls below 0.1. This stabilizes training by starting as a stable, shallow unrolled network and progressively refining to a precise equilibrium state.

Architectural Layout

The model implements a Sandwich Architecture consisting of three sequential stages:

 ┌──────────────────────────────┐
 │ Token & Pos Embeddings │
 └──────────────┬───────────────┘
 ▼
 ┌──────────────────────────────┐
 │ Fixed Encoder Stack (M) │ <-- Standard Transformers
 └──────────────┬───────────────┘
 ▼
 ┌──────────────────────────────┐
 │ Recurrent Core Stack │ <-- Mamba or Transformer
 │ (Iterated until Eq/Halt) │ (Width: n-fixed-core)
 └──────────────┬───────────────┘
 ▼
 ┌──────────────────────────────┐
 │ Fixed Decoder Stack (K) │ <-- Standard Transformers
 └──────────────┬───────────────┘
 ▼
 ┌──────────────────────────────┐
 │ LM Head / Output Logits │
 └──────────────────────────────┘

Setup & Installation

The project uses uv for python virtual environments and package management.

1. Install Dependencies

uv sync

2. Run Verification Tests

Verify the code correctness, shape constraints, and backward gradient passes:

uv run python test_dataset.py

Usage Instructions

1. Training the Model

Train the network on-the-fly with customizable architectures.

• Fast DEQ Training (Transformer Core + 1st Order Taylor Backprop + Compilation):

  uv run python train.py --model-type deq --core-type transformer --taylor-order 1 --n-fixed-core 2 --spectral-norm --core-lr-factor 0.5 --weight-decay 1e-2 --compile
  

• Standard ACT Training:

  uv run python train.py --model-type act --core-type transformer --max-steps 8 --n-fixed-core 1
  

Key Training CLI Arguments:

• --model-type: Select act or deq.
• --core-type: Select transformer or mamba.
• --n-fixed-core: Number of sequential blocks inside the recurrent core.
• --taylor-order: Adjoint solver order for DEQ backprop (-1 for full solver, 0 for Identity, 1 for 1st order Taylor).
• --epsilon: Tolerances/halting thresholds.
• --spectral-norm: Apply spectral normalization to the core.
• --compile: Compile the model via PyTorch Inductor (defaults to True on Linux).
• --no-compile: Disable compilation (overrides --compile, useful if Triton is unavailable).

2. Interactive Calculator Mode

Test the trained checkpoints in an interactive command-line interface:

uv run python interactive.py --checkpoint model_best.pt --model-type deq --core-type transformer

Codebase Map

• dataset.py: On-the-fly rational math generation, tokenization, and batch tensor collation.
• model.py: Core layers (RMSNorm, TransformerBlock, MambaBlock), custom autograd DEQFunction, and the UniversalTransformerWithACT class.
• train.py: Training loop, evaluation routines, curriculum scaling, and TensorFloat-32 config.
• interactive.py: Interactive CLI generator.
• test_dataset.py: Automated PyTest-compatible unit tests.