LDT: Layer-Decomposition Training Makes Networks More Generalizable (ICLR2026)

[Paper]

Zaizuo Tang¹, Zongqi Yang¹, Yubin Yang¹,

¹State Key Laboratory for Novel Software Technology, Nanjing University, Nanjing, China

Abstract

Domain generalization methods can effectively enhance network performance on test samples with unknown distributions by isolating gradients between unstable and stable parameters. However, existing methods employ relatively coarse-grained partitioning of stable versus unstable parameters, leading to misclassified unstable parameters that degrade network feature processing capabilities. We first provide a theoretical analysis of gradient perturbations caused by unstable parameters. Based on this foundation, we propose Layer-Decomposition Training (LDT), which conducts fine-grained layer-wise partitioning guided by parameter instability levels, substantially improving parameter update stability. Furthermore, to address gradient amplitude disparities within stable layers and unstable layers respectively, we introduce a Dynamic Parameter Update (DPU) strategy that adaptively determines layer-specific update coefficients according to gradient variations, optimizing feature learning efficiency. Extensive experiments across diverse tasks (super-resolution, classification) and architectures (Transformer, Mamba, CNN) demonstrate LDT's superior generalization capability. Our code is available at https://github.com/ZaizuoTang/LDT.

Train on SR

1. Download the DRealSR dataset.

   [DRealSR Dataset]

2. Download the pre-trained weights.

   [MambaIR V1 classicSRx4.pth]

3. Start training:

   python Train_pip.py

Test on SR

python basicsr/test.py -opt test_MambaIR_SR_x4.yml

License

This project is released under the Apache 2.0 license.

Acknowledgement

This code is based on BasicSR and MambaIR. Thanks for their awesome work.

Contact

If you have any questions, feel free to approach me at tangzz@smail.nju.edu.cn