9 papers
Learning to Extrapolate in Depth with Looped Transformers
Huitao Yang, Jianhao Huang, Zixuan Wang +1
Looped transformers reuse a single block across iterations, raising a natural question: how far can a block trained at one depth be safely extrapolated to greater depths? We study…
Learning What to Remember: Test-Time Training via Context Distillation
Zixuan Wang, Xingyu Dang, Rui-Jie Zhu +4
Effective long-context modeling is not merely about retaining more of the past, but about preserving the information that may prove relevant later. Test-time training (TTT) is an a…
Learning Compositional Functions with Transformers from Easy-to-Hard Data
Zixuan Wang, Eshaan Nichani, Alberto Bietti +4
Transformer-based language models have demonstrated impressive capabilities across a range of complex reasoning tasks. Prior theoretical work exploring the expressive power of tran…
Beyond Words: Augmenting Discriminative Richness via Diffusions in Unsupervised Prompt Learning
Hairui Ren, Fan Tang, He Zhao +3
Fine-tuning vision-language models (VLMs) with large amounts of unlabeled data has recently garnered significant interest. However, a key challenge remains the lack of high-quality…
What Makes a Reward Model a Good Teacher? An Optimization Perspective
Noam Razin, Zixuan Wang, Hubert Strauss +3
The success of Reinforcement Learning from Human Feedback (RLHF) critically depends on the quality of the reward model. However, while this quality is primarily evaluated through a…
Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought
Jianhao Huang, Zixuan Wang, Jason D. Lee
Chain of Thought (CoT) prompting has been shown to significantly improve the performance of large language models (LLMs), particularly in arithmetic and reasoning tasks, by instruc…
SeMi: When Imbalanced Semi-Supervised Learning Meets Mining Hard Examples
Yin Wang, Zixuan Wang, Hao Lu +7
Semi-Supervised Learning (SSL) can leverage abundant unlabeled data to boost model performance. However, the class-imbalanced data distribution in real-world scenarios poses great…
Learning and Transferring Sparse Contextual Bigrams with Linear Transformers
Yunwei Ren, Zixuan Wang, Jason D. Lee
Transformers have excelled in natural language modeling and one reason behind this success is their exceptional ability to combine contextual informal and global knowledge. However…
Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot
Zixuan Wang, Stanley Wei, Daniel Hsu +1
The transformer architecture has prevailed in various deep learning settings due to its exceptional capabilities to select and compose structural information. Motivated by these ca…