LaCT

Test-Time Training Done Right

<a href="https://tianyuanzhang.com/">Tianyuan Zhang</a><sup>1</sup>,
<a href="https://sai-bi.github.io/">Sai Bi</a><sup>2</sup>,
<a href="https://yiconghong.me/">Yicong Hong</a><sup>2</sup>,
<a href="https://kai-46.github.io/website/">Kai Zhang</a><sup>2</sup>,
<a href="https://scholar.google.com/citations?user=NLxrmYQAAAAJ">Fujun Luan</a><sup>2</sup>,
<a href="https://sustcsonglin.github.io/">Songlin Yang</a><sup>1</sup>,
<a href="http://www.kalyans.org/">Kalyan Sunkavalli</a><sup>2</sup>,
<a href="https://billf.mit.edu/">William T. Freeman</a><sup>1</sup>,
<a href="https://www.cs.unc.edu/~airsplay/">Hao Tan</a><sup>2</sup>



<sup>1</sup>MIT <sup>2</sup>Adobe Research

Paper | Website | Models (HuggingFace)

Minimal Implementations

We provide minimal implementations for a LaCT layer in minimal_implementations/. This implementation serves as a starting point for understanding, modifying, and creating your own version of LaCT.

News

November 18: Added Triton kernels for the fused TTT layer (see lact_llm/lact_model/ttt_operation_fused_kernel.py for implementation, with a updated moddel config in configs/760M_lact_swiglu_nh4_fwlow_rank_momentum_fused_kernel.json). This reduces training memory consumption. The triton kernel fuse several matmul with epilogues to reduce read and writes in global mememory.

Plans

• ~~Release Language model code.~~ [Done] (will also release inside this great repo: https://github.com/fla-org/flame)
• ~~Release novel view synthesis training code and model by June 12.~~ [Done]
• ~~Release video model finetuning code (build on top of Wan T2V) by June 24.~~ [Done]

Citations

If you find this codebase useful for your research, please kindly cite our paper:

@article{zhang2025test,
  title={Test-time training done right},
  author={Zhang, Tianyuan and Bi, Sai and Hong, Yicong and Zhang, Kai and Luan, Fujun and Yang, Songlin and Sunkavalli, Kalyan and Freeman, William T and Tan, Hao},
  journal={arXiv preprint arXiv:2505.23884},
  year={2025}
}