ReThinkLab-SJTU

Founded by Prof. Junchi Yan since 2018, ReThinkLab's vision is to perform recurrent thinking to change the world through both its technology and talents.

[Yang Li @ ICML @LeYangco]

Our work “Attention Illuminates LLM Reasoning” will be presented at ICML 2026. Happy to connect and discuss with anyone interested in LLM reasoning, interpretability, and RL4LLM.

This work investigates how attention patterns can reveal the internal reasoning structure of LLMs. We find that attention is not merely a by-product of generation, but often reflects a structured reasoning rhythm: local chunking, global semantic anchors, and a recurring preplan-anchor coupling pattern.

Based on these observations, we further propose an attention-guided reward redistribution method for RL, aligning optimization with the model’s intrinsic reasoning structure. This leads to more fine-grained, interpretable, and efficient reinforcement learning for reasoning models.

Paper: https://t.co/zDMoZ8SLgK

[Yang Li @ ICML @LeYangco]

How does reasoning actually flow inside LLMs?

In our ICML 2026 paper, How Does Reasoning Flow? Tracing Attention-Induced Information Flow for Targeted RL in LLMs (https://t.co/53SaWFd7TX), we trace attention-induced information flow to reveal the “main roads” of reasoning inside language models.

Instead of treating a reasoning trace as a flat sequence, we view it as an information-flow network: problem facts, intermediate conclusions, variables, symbols, and step structures are read, integrated, forwarded, and eventually merged into the final answer.

This lets us ask a sharper question:

Which tokens are truly on the effective path toward the answer?

We find that reasoning flow is far from uniform. It shows rhythmic “gathering → redistribution” patterns, with high-flow tokens acting as critical hubs — often step boundaries, repeated variables, key numbers, or operators.

When these hubs are disrupted, models are much more likely to derail.

Building on this, we use reasoning-flow attribution for more targeted RL: instead of spreading reward uniformly across the whole response, we route stronger training signals to the key tokens that actually support answer formation.

We’ll be presenting this work at ICML 2026. Come find us at ICML 2026 if you’re interested！

[ReThinkLab-SJTU @Rethinker135365]

🚀 Happy to present our new work on LLM Scaling Laws: JTok!

We show that token-indexed parameters can serve as a novel, orthogonal scaling axis that decouples model capacity from FLOPs. By modulating Transformer layers with JTok, we achieve comparable model quality with 35% less compute relative to vanilla MoE architectures!

⚙️ JTok & JTok-M Architecture

(1) Local Modulation: JTok introduces an auxiliary embedding table at each transformer layer, retrieving token-specific vectors to modulate the backbone through element-wise operations, incurring negligible FLOPs overhead.

(2) Dynamic Mixture (JTok-M): To further scale this embedding table, JTok-M introduces an embeeding pool and uses a lightweight router to select a sparse Top-K mixture per token.

(3) System Efficiency: Retrieval overlaps with computation. CPU offloading ensures zero extra GPU memory footprint with <7.3% latency increase.

🧠 Key Scaling & Performance Results

(1) 35% Compute Savings: Rigorous IsoFLOPs analysis confirms that JTok-M fundamentally shifts the quality-compute Pareto frontier. It achieves comparable model quality while saving 35% compute relative to vanilla MoE models.

(2) Massive Downstream Gains: We validated this on backbones up to 61B total parameters (17B backbone + 44B embedding). On a large-scale 17B MoE backbone, JTok-M delivers substantial improvements, including +4.1 on MMLU, +8.3 on ARC, and +8.9 on CEval.

(3) Predictable Power-Law Scaling: Validation loss exhibits a log-linear trend with the number of token-indexed parameters. This establishes token-indexed parameters as a highly efficient and scalable dimension alongside traditional dense and sparse scaling.

🎉 Easter Egg: What does "JTok" mean? Technically, it stands for "Joint-Token". But for us SJTUers, "Joint" sounds like "Jiao Tong" (交通), hiding our inside joke: "JT (交通) OK!" 🎓 Following the "jAccount" naming tradition, this is our tribute to Shanghai Jiao Tong University's upcoming 130th anniversary! 🎂❤️

arXiv Link: https://t.co/8clTCMCjm7

[Xiaohan Qin @qinxiaohan94414]

🚀 No reference model, yet better token selection?
Introducing ssToken — for smarter &amp; cheaper SFT fine-tuning!
It improves token-level data selection without extra models, combining learnability + semantics.

📄 HF: https://t.co/Lf7PlG4CXQ
📘 arXiv: https://t.co/VWQpZeW8ad

[Zhanpeng Zhou @zhanpeng_zhou]

Two papers get accepted by #ICML2025 🥳🥳

[1/2] We discover that different blocks in Transformers exhibit notable disparity in Sharpness. Then we propose Blockwise LR, accelerating large language model (LLM) pre-training (~2x speedup).

https://t.co/MZ3wWujL0I

[Xiangdong Zhang @ ICML26 @aHapBean]

🚀 Can video understanding boost video generation? Introducing VideoREPA (NeurIPS’25)

State-of-the-art text-to-video models generate visually stunning results—but still violate basic physics (floating objects, collisions ignored), limiting their reliability as world models.

[Yang Li @ ICML @LeYangco]

🚀 Happy to present our new work on LLM reasoning!
We show that: (1) Attention is a structured map of the model's reasoning logic, uncovering a preplan-and-anchor reasoning rhythm. (2) Aligning RL objectives with the model's intrinsic attention rhythm yields more transparent, fine-grained, and efficient optimization.

🧠 Key Reasoning Patterns in Attention
(1) Local Chunking: Near-diagonal sawtooth patterns indicate dense intra-chunk processing. At chunk boundaries, the model performs long-range context retrieval (often with higher entropy), which guides subsequent generation.
(2) Global Anchor Planning: Sparse, high-influence anchor tokens exert broad control over later tokens. Perturbing these anchors significantly disrupts downstream reasoning.
(3) Preplan-Anchor Coupling: A stable temporal rhythm emerges: the model first emits a "preplan" token, then anchors a core semantic node, repeatedly structuring the reasoning trajectory.

⚙️ RL Innovation
We introduce a dynamic reward redistribution mechanism guided by attention-derived reasoning structure:
(1) Preplan Guidance: Boosts tokens that guide local chunks and enable long-range referencing.
(2) Anchor Enhancement: Prioritizes optimization of globally influential semantic anchors.
(3) Coupling Alignment: Reinforces the temporal coordination between preplans and anchors to solidify structured reasoning.

HuggingFace Link: https://t.co/6JsKknOpjV
arXiv Link: https://t.co/zDMoZ8Sdrc

#LLMs #artificial_intelligence #RL4LLM