Jiaju Wu

MATHEMATICS & MACHINE LEARNING

Jiaju Wu

I am a mathematics undergraduate at Peking University. I study how Transformers compute and generalize, with a focus on the mechanisms of looped Transformers.

Jiaju Wu presenting a project in a classroom
Latest paper · 2026Submitted to ICLR 2027

Shared Weights, Selected Computations

How Looped Transformers Route What Each Loop Does

Jiaju Wu, Yi Hu, Muhan Zhang

We study how a changing hidden state selects different computations in a looped Transformer with fixed weights. Graph-walk experiments and attention interventions test how this selection works.

Figure 1 · Paper overview

Research

QUESTION

What does each loop compute?

Shared weights do not require successive loops to perform the same operation.

FINDING

The entering state controls computation.

In the studied graph-walk models, changing the entering hidden state redirects a frozen loop. Attention interventions test the route of this control.

SCOPE

Control has limits.

Training changes which transitions can be selected. The evidence is task- and model-specific; longer compositions remain a limitation.

Experiment index ↗

Archived results support figure and statistic reproduction; full training requires additional dependencies and checkpoints.

Earlier research note · S₅ circuits
Three observed representation circuit families in S5 multiplication
Circuit summary reconstructed from the research note.

Research note

Same Answers, Different Circuits

I studied how Transformers compute multiplication in S₅. Representation analysis and activation interventions identified different internal implementations of the same function.

The experiments identified three families of product representations. Models computed these representations through different MLP, attention, and residual paths. Optimizer preferences varied across seeds.

The study uses a small set of S₅ models. It does not establish a deterministic mapping from optimizer to circuit.

Mechanistic interpretability · Group representations

Honors

  • 2025Lingjun Linghang University Scholarship ↗
  • 2025Chinese Mathematics Competitions · Beijing Division, Mathematics A · First Prize
  • 2025China Undergraduate Mathematical Contest in Modeling · Second Prize
  • 2022 / 2023Chinese High School Mathematics League · First Prize
  • 2023Chinese Mathematical Olympiad · Gold Medal ↗
  • 2023Southeast Mathematical Olympiad · Gold Medal
  • 2024Chinese Mathematics Competitions · Beijing Division, Mathematics A · Second Prize

Education

AUG 2024 — PRESENT

Peking University

School of Mathematical Sciences · Undergraduate

GPA: 3.79 / 4.00

Core coursework

Mathematical Analysis I–III; Geometry; Advanced Algebra I–II; Probability Theory; Mathematical Statistics; Data Structures and Algorithms; Introduction to Mathematical Machine Learning; Ordinary Differential Equations; Deep Learning Theory; Multi-Agent Foundations; Applied Stochastic Processes.

2022 — 2024

Weifang Beichen High School

Early admission to the PKU Mathematics Talent Program

Selected Projects

Course project

Learning Dynamics & Schedule Transfer

Predicting WSD loss curves from a cosine training curve, with a frozen MPL baseline and source-only calibration.

The reported correction improves all 15 same-scale targets. Evaluation uses existing public curves, not new model training.

Team course project

Grokking on Algorithmic Tasks

A joint study with Kairui Li and Kehan Huang of delayed generalization on modular arithmetic and group operations.

Shared training and evaluation tools for Transformers, MLPs, LSTMs, and GRUs; code and a course report are available.

Campus software

A Website for the Academic Community

A home for the SMS Student Union Academic Department's problems, course resources, and events.

React and TypeScript, with mathematical typesetting and configuration-based content updates.

More projects & teaching
Causal Transformer for High-Frequency Return Prediction

Lingjun Quant Challenge project. Built a causal Transformer for A-share high-frequency microstructure data with 500 stocks, 239 intraday minutes, and 384 features, predicting ten-minute-ahead returns with strict time-split validation and leakage-aware preprocessing.

Project notes ↗
LLM Math Text Detection and Attribution

Midterm project for Mathematical Modeling. Built a source-attribution system for mathematical solutions from humans and major LLMs including DeepSeek, GLM, Kimi, and Qwen, combining feature-based baselines with neural text classification.

Project notes ↗
MultiagentFinal: Intent-Grounded Cooperative Sokoban

Final project for Multi-Agent Foundations. Built StrictCoop-Sokoban and studied intent-grounded recurrent communication for partially observable cooperative multi-agent reinforcement learning.

Project notes ↗
Tuvalon: Avalon Agent Battle Platform

Teaching-assistant project design for Data Structures and Algorithms B. Developed and maintained a Flask-based Avalon agent battle platform where submitted Python agents compete under partial observability, communication constraints, referee-controlled game phases, automatic matching, ELO ranking, and replay-oriented logs.

Project notes ↗