LaCT

Test-Time Training Done Right

Tianyuan Zhang¹, Sai Bi², Yicong Hong², Kai Zhang², Fujun Luan², Songlin Yang¹, Kalyan Sunkavalli², William T. Freeman¹, Hao Tan^2
1MIT ²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}
}