5 papers
Finding the Minimal Parameter Budget for Implicit Reasoning: A Data Complexity Driven Scaling Law for Language Models
Xinyi Wang, Shawn Tan, Shenbo Xu +4
Reasoning is a core capability of language models (LMs), yet it remains unclear how much model capacity is necessary to support reasoning during pretraining. In this work, we study…
MRNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling
Mayank Mishra, Shawn Tan, Ion Stoica +2
Transformers are highly parallel but are limited to computations in the TC complexity class, excluding tasks such as entity tracking and code execution that provably require gr…
PaTH Attention: Position Encoding via Accumulating Householder Transformations
Songlin Yang, Yikang Shen, Kaiyue Wen +5
The attention mechanism is a core primitive in modern large language models (LLMs) and AI more broadly. Since attention by itself is permutation-invariant, position encoding is ess…
Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study
Shawn Tan, Songlin Yang, Aaron Courville +2
The self-attention mechanism traditionally relies on the softmax operator, necessitating positional embeddings like RoPE, or position biases to account for token order. But current…
Scattered Mixture-of-Experts Implementation
Shawn Tan, Yikang Shen, Rameswar Panda +1
We present ScatterMoE, an implementation of Sparse Mixture-of-Experts (SMoE) on GPUs. ScatterMoE builds upon existing implementations, and overcoming some of the limitations to imp…