Factorizing Knowledge in Neural Networks (2207.03337v2)

Abstract: In this paper, we explore a novel and ambitious knowledge-transfer task, termed Knowledge Factorization~(KF). The core idea of KF lies in the modularization and assemblability of knowledge: given a pretrained network model as input, KF aims to decompose it into several factor networks, each of which handles only a dedicated task and maintains task-specific knowledge factorized from the source network. Such factor networks are task-wise disentangled and can be directly assembled, without any fine-tuning, to produce the more competent combined-task networks. In other words, the factor networks serve as Lego-brick-like building blocks, allowing us to construct customized networks in a plug-and-play manner. Specifically, each factor network comprises two modules, a common-knowledge module that is task-agnostic and shared by all factor networks, alongside with a task-specific module dedicated to the factor network itself. We introduce an information-theoretic objective, InfoMax-Bottleneck~(IMB), to carry out KF by optimizing the mutual information between the learned representations and input. Experiments across various benchmarks demonstrate that, the derived factor networks yield gratifying performances on not only the dedicated tasks but also disentanglement, while enjoying much better interpretability and modularity. Moreover, the learned common-knowledge representations give rise to impressive results on transfer learning. Our code is available at https://github.com/Adamdad/KnowledgeFactor.

Factorizing Knowledge in Neural Networks

The paper "Factorizing Knowledge in Neural Networks" presents a novel approach to knowledge transfer in deep neural networks, termed Knowledge Factorization (KF). The primary objective of KF is to decompose a pretrained network into a set of task-specific sub-networks, called factor networks, which can be reassembled to form comprehensive networks for multiple tasks without additional training. This modular and assemblable approach is akin to constructing systems from Lego-like building blocks.

Core Concept and Methodology

The authors address the limitations of conventional Knowledge Distillation (KD) methods by introducing a model reusability scheme that emphasizes modularization and interpretability. Unlike KD, where a student model learns a condensed version of a teacher's capabilities, KF aims to achieve both structural and representational disentanglement.

Structural Factorization is achieved by dividing each factor network into a Common Knowledge Network (CKN) and a Task-Specific Network (TSN). The CKN is responsible for capturing shared, task-agnostic information across the tasks, while the TSN handles task-specific knowledge. This setup allows each factor network to specialize in one task using both task-agnostic and task-specific features, thereby achieving architectural disentanglement.

Representation Factorization is driven by the InfoMax Bottleneck (IMB) objective, which enforces statistical independence of task-specific features. IMB maximizes mutual information between the input data and common features, while minimizing it for task-specific features. This ensures that task features only retain information pertinent to their designated task, enhancing interpretability and robustness.

Experimental Results

The performance of the KF approach was evaluated on synthetic datasets like dSprites and Shape3D, as well as real-world classification and multi-task datasets such as CIFAR-10 and ImageNet1K. The factor networks demonstrated superior performance compared to both single-task and multi-task baselines and retained high capacity for knowledge transfer, especially when reassembling networks for multiple tasks.

On synthetic datasets, KF outperformed KD and baseline models in terms of ROC-AUC scores, indicating its efficacy in classification tasks. For real-world datasets, factor networks outperformed their counterparts by showing a consistent improvement in classification accuracy, with significant gains observed in multi-task settings where task-specific networks could be assembled from CKN and relevant TSNs.

Disentanglement Metrics: The investigation included disentanglement metrics and representation similarity using CKA. The results showed that factor networks achieved better disentanglement, with higher DCI and MIG scores, indicating the successful capture of independent variables across tasks.

Implications and Future Work

The proposed KF offers promising applications in edge computing and other resource-constrained environments where adaptability and interpretability are critical. By structurally and representationally factorizing knowledge, this approach facilitates scalable multi-task learning and offers a pathway to efficient model reusability.

Future research could focus on expanding KF to other domains, like natural language processing or reinforcement learning, where task interdependencies are more complex. Additionally, integrating more advanced information-theoretic methods may yield further improvements in disentanglement and transferability across diverse tasks and datasets.

Conclusion

The Knowledge Factorization framework stands out as an innovative technique in the landscape of neural network optimization and knowledge transfer. By using modular networks and information-theoretic objectives, it provides a flexible and interpretable method for building and adapting multi-task networks. This approach aligns well with future demands for modular AI systems that can dynamically integrate new tasks without re-training from scratch, making it a valuable contribution to the field of deep learning and beyond.