Georgia Institute of Technology · GTISC Lab · Atlanta
I am a Computer Science Ph.D. student at Georgia Tech, advised by Prof. Wenke Lee in the GTISC Lab. My research focuses on understanding and improving reasoning in language models. I study useful structure in model computation, reasoning trajectories, and interaction data, and use these insights to guide model design, post-training, and inference algorithms.
Before joining the Ph.D. program, I earned my B.Eng. in Communication Engineering (Artificial Intelligence Track) from Wuhan University.
I study how intermediate representations and generation dynamics affect model computation, from information flow across diffusion steps and state decay in long contexts to progress-conditioned reasoning and redundancy-aware decoding.
I study how interaction data and task structure can guide learning and evaluation: evidence dependencies supervise observation decisions, task feedback refines reusable source understanding, and task dependencies guide executable benchmark construction.
Selected publications
2026
Standard discrete diffusion models discard continuous representations between denoising steps. MetaState learns a recurrent pathway around a frozen backbone to carry this information forward, improving mathematical reasoning and code generation.
@misc{metastate2026,
title = {MetaState: Persistent Working Memory Enhances Reasoning in Discrete Diffusion Language Models},
author = {Kejing Xia and Mingzhe Li and Lixuan Wei and Zhenbang Du and Xiangchi Yuan and Dachuan Shi and Qirui Jin and Wenke Lee},
year = {2026},
eprint = {2603.01331},
archivePrefix = {arXiv},
primaryClass = {cs.CL},
url = {https://arxiv.org/abs/2603.01331}
}
R²-dLLM identifies spatial prediction patterns and cross-step stability in diffusion trajectories, using them to guide training-free decoding and redundancy-aware supervised fine-tuning. It reduces decoding steps while maintaining competitive generation quality.
@misc{r2dllm2026,
title = {{R$^2$-dLLM}: Accelerating Diffusion Large Language Models via Spatio-Temporal Redundancy Reduction},
author = {Zhenbang Du and Kejing Xia and Xinrui Zhong and Yonggan Fu and Nicolai Oswald and Binfei Ji and Brucek Khailany and Pavlo Molchanov and Yingyan (Celine) Lin},
year = {2026},
eprint = {2604.18995},
archivePrefix = {arXiv},
primaryClass = {cs.CL},
url = {https://arxiv.org/abs/2604.18995}
}
SourceLearn turns repeated use of a knowledge source into reusable understanding. Self-directed study and task feedback reveal gaps in an external source model. Updates are reconstructed from the authoritative source.
@misc{sourcelearn2026,
title = {From Knowledge Access to Source Learning: Developing Source-Specific Competence},
author = {Lucheng Fu and Kejing Xia and Yiyang Wang and Yiqiao Jin and Jinjin He and Xiyuan Yang and Haoxin Liu and Ye Yu and Haibo Jin and Yijia Xiao and Wenke Lee and B. Aditya Prakash and Haohan Wang},
year = {2026},
eprint = {2610.02150},
archivePrefix = {arXiv},
primaryClass = {cs.CL},
url = {https://arxiv.org/abs/2610.02150}
}
2025
D²GS turns predicted depth into dense geometric supervision for LiDAR-free urban reconstruction. Diffusion priors refine depth rendered from Gaussians, which in turn constrains geometry, alongside progressive pruning and road-specific constraints.
@inproceedings{NEURIPS2025_524ff06d,
author = {Xia, Kejing and Jia, Jidong and Jin, Ke and BAI, Yucai and Sun, Li and Tao, Dacheng and Zhang, Youjian},
booktitle = {Advances in Neural Information Processing Systems},
doi = {10.52202/085713-1905},
editor = {D. Belgrave and C. Zhang and H. Lin and R. Pascanu and P. Koniusz and M. Ghassemi and N. Chen},
pages = {56899--56927},
publisher = {Curran Associates, Inc.},
title = {D\^{}2GS: Dense Depth Regularization for LiDAR-free Urban Scene Reconstruction},
url = {https://proceedings.neurips.cc/paper_files/paper/2025/file/524ff06d5375d76e93d8490e654a81bf-Paper-Conference.pdf},
volume = {38, Main Conference},
year = {2025}
}
EDITOR starts from an image caption, optimizes contextual text embeddings with image-reconstruction feedback, and converts them into readable prompts. This improves visual reconstruction while retaining interpretable text.
@misc{editor2025,
title = {EDITOR: Effective and Interpretable Prompt Inversion for Text-to-Image Diffusion Models},
author = {Mingzhe Li and Kejing Xia and Gehao Zhang and Zhenting Wang and Guanhong Tao and Siqi Pan and Juan Zhai and Shiqing Ma},
year = {2025},
eprint = {2506.03067},
archivePrefix = {arXiv},
primaryClass = {cs.CV},
url = {https://arxiv.org/abs/2506.03067}
}
LongMamba traces long-context failure to cumulative state decay in Mamba's global channels. Targeted token filtering changes which updates accumulate in these channels, enlarging their receptive fields without additional training.
@inproceedings{ICLR2025_ab5d50d2,
author = {Ye, Zhifan and Xia, Kejing and Fu, Yonggan and Dong, Xin and Hong, Jihoon and Yuan, Xiangchi and Diao, Shizhe and Kautz, Jan and Molchanov, Pavlo and Lin, Yingyan Celine},
booktitle = {International Conference on Learning Representations},
editor = {Y. Yue and A. Garg and N. Peng and F. Sha and R. Yu},
pages = {68486--68504},
title = {LongMamba: Enhancing Mamba\textquotesingle s Long-Context Capabilities via Training-Free Receptive Field Enlargement},
url = {https://proceedings.iclr.cc/paper_files/paper/2025/file/ab5d50d269e52f8eed497062311ff173-Paper-Conference.pdf},
volume = {2025},
year = {2025}
}
By examining attention patterns in state space models, LAMB guides token filtering to preserve critical information, strengthening long-context understanding in pure and hybrid models without training.
@inproceedings{ye-etal-2025-lamb,
title = "{LAMB}: A Training-Free Method to Enhance the Long-Context Understanding of {SSM}s via Attention-Guided Token Filtering",
author = "Ye, Zhifan and
Wang, Zheng and
Xia, Kejing and
Hong, Jihoon and
Li, Leshu and
Whalen, Lexington and
Wan, Cheng and
Fu, Yonggan and
Lin, Yingyan Celine and
Kundu, Souvik",
editor = "Che, Wanxiang and
Nabende, Joyce and
Shutova, Ekaterina and
Pilehvar, Mohammad Taher",
booktitle = "Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)",
month = jul,
year = "2025",
address = "Vienna, Austria",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2025.acl-short.96/",
doi = "10.18653/v1/2025.acl-short.96",
pages = "1200--1209",
ISBN = "979-8-89176-252-7"
}
LookAgain-CUA turns evidence dependencies in GUI demonstrations into supervision for when to observe again, then refines the policy with cost-aware RL. It reduces model calls and estimated inference cost on OSWorld at comparable task success.
RoT uses relative position within reasoning traces as intermediate supervision. An online predictor estimates this progress from hidden states and conditions decoder computation, with the resulting framework trained through reinforcement learning.
GEvol uses task-dependency graphs to compose executable benchmark edits while preserving original goals and verifiers. Joint validation and independent execution feedback guide the construction of more challenging terminal-agent tasks.
Ph.D. student
Jan. 2026 – Present
MetaState adds persistent working-memory augmentation for reasoning across diffusion steps. RoT is an RL-trained latent reasoning framework that modulates decoder layers with an uncertainty-aware estimate of trace-relative reasoning progress. SourceLearn builds reusable, source-specific competence for LLM agents through an external source model.
Undergraduate to Ph.D. student
May 2024 – Dec. 2025
Projects study efficient diffusion decoding and long-context memory retention in state space models, including redundancy-aware decoding, soft token interpolation, and token filtering.
Undergraduate student
Oct. 2023 – Apr. 2024
BRS-E²NeRF uses an event-and-RGB pipeline to jointly address rolling-shutter distortion and motion blur for sharp 3D reconstruction and novel-view synthesis.
Research Intern
May 2026 – Aug. 2026
LookAgain-CUA learns observation boundaries and executes several GUI actions per model call, using multi-action supervision and cost-aware RL. On OSWorld, it reduces model calls and estimated inference cost at comparable task success.
Research Intern
Jul. 2024 – May 2025
D²GS co-optimizes diffusion-based depth generation and 3D Gaussian Splatting for LiDAR-free urban reconstruction, with road-specific geometric constraints and benchmarking on Waymo.
Aug. 2025 – Present
Ph.D. in Computer Science · In progress
Sept. 2021 – Jun. 2025
B.Eng. in Communication Engineering (Artificial Intelligence Track)
First-Class Academic Scholarship (2023–2024)
I studied at the No. 1 Middle School Affiliated to Central China Normal University.