The author developed a small transformer model from scratch that achieved 44% on the ARC-AGI-1 benchmark, matching the performance of TRM/HRM models at a significantly lower cost of approximately 67 cents using an NVIDIA 5090. The approach utilizes a supervised training regime where the loss function is applied only to output tokens, rather than both input and output tokens. To facilitate cross-task learning and robust representations, the architecture incorporates per-task additive embeddings and 3D RoPE (Rotary Positional Embedding). Key architectural upgrades from previous iterations include the implementation of SwiGLU gated feed-forward networks, RMSNorm, and an 8-layer depth. Efficiency gains were achieved by replacing AdamW with the NorMuon optimizer and utilizing Flash Attention with variable-length training and flex-attention kernels for inference. While the model relies on color and dihedral permutations for data augmentation, the researcher emphasizes that the performance gains stem primarily from improved representation and modern architectural components rather than synthetic data or heavy inductive biases. Ablation studies indicate that 3D RoPE and task embeddings are critical, as removing them drops performance from 44% to roughly 24%.
Read original
hackernews