PaperlayerView exact revision on arXiv

Source-linked AI summary

Factorizing Knowledge in Neural Networks

Xingyi Yang, Jingwen Ye, Xinchao Wang

arXiv:2207.03337v2cs.CVcs.AIcs.LG

TL;DR

Conventional knowledge distillation is opaque and assumes homogeneous teacher and student tasks, limiting its flexibility and scalability. The paper proposes Knowledge Factorization, which separates shared and task-specific knowledge into assemblable factor networks using an information-theoretic objective. Across benchmarks, the authors report strong performance, disentanglement, modularity, and transferability.

  • Problem

    Conventional knowledge distillation obscures the knowledge-transfer process and assumes homogeneous teacher and student tasks, while learning many task-specific students requires repeated distillation.

    Source: §1 · p. 2

  • Method

    Knowledge Factorization decomposes a pretrained teacher into task-disentangled factor networks with shared common-knowledge and task-specific modules, using an InfoMax Bottleneck objective to separate their representations.

    Source: §1 · pp. 3–4

  • Results

    Across various benchmarks, the method achieves strong and robust performance with disentanglement, better modularity, and transferability.

    Source: §1 · p. 4, §5 · p. 14

  • Takeaways & Limitations

    Factor networks can operate independently or be assembled into multi-task networks, and common-knowledge representations support transfer to unseen downstream tasks.

    Source: §1 · p. 3, §5 · p. 14

    The variational task-feature model assumes a zero-mean, unit-variance Gaussian prior and a Gaussian distribution for p(tj|x).

    Source: §3.4 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames conventional knowledge distillation as opaque and poorly suited to efficiently producing students for different tasks. It proposes Knowledge Factorization (KF), which splits a pretrained model into task-disentangled factors that can be assembled, and reports strong performance and transferability across benchmarks.

1 Introduction

  • Conventional knowledge distillation transfers a pretrained teacher’s expertise to a smaller student, but its knowledge flow is opaque and its task setup is restrictive.Distilling many task-specific students separately can become computationally intractable as the number of task combinations grows.Source: §1 · p. 2
  • KF decomposes a pretrained teacher into task-wise disentangled factor networks that can be assembled without fine-tuning into a multi-task network.The factors are presented as plug-and-play building blocks for customized networks.Source: §1 · p. 3
  • The authors present KF as a way to make knowledge transfer more interpretable and to let owners selectively transfer task knowledge without leaking knowledge of other tasks.The introduction attributes this protection to the factors’ disentangled, task-specific knowledge.Source: §1 · pp. 3–4
  • For n knowledge types, the introduction contrasts sub-KD’s 2^n full-sized models with KF’s one full-sized common model and n smaller task models.The smaller task models have s parameters, where s ≪ S, compared with S parameters per sub-KD model.Source: §1 · p. 3
  • KF combines a shared task-agnostic common-knowledge network with a task-specific network in each factor network.The common module captures task commonality, while the task-specific module handles its dedicated information.Source: §1 · p. 3
  • InfoMax Bottleneck (IMB) factorizes shared and task-level representations by maximizing input–common-feature mutual information and minimizing data–task mutual information.The paper derives a variational lower bound to optimize the objective.Source: §1 · pp. 3–4
  • KF achieves strong performance and disentanglement across benchmarks, with better modularity and transferability.The introduction also reports architecture-level and representation-level disentanglement and transfer to unseen downstream tasks.Source: §1 · p. 4

2 Related work

The related work contrasts whole-model knowledge transfer with task-wise factorization, and positions the approach alongside disentangled representation learning, information-theoretic learning, and multi-task learning.

Knowledge Distillation

  • Knowledge distillation transfers knowledge from a model or ensemble to a student, whereas this work factorizes a multi-talented teacher into task-disentangled networks.The paper describes conventional distillation as disseminating knowledge as a whole.Source: §2 · p. 4

Disentangled representation learning

  • Prior disentanglement methods commonly use adversarial learning or variational auto-encoders; this work targets task-agnostic and task-related representations through mutual information.Source: §2 · p. 4

InfoMax principle and Information bottleneck

  • The paper combines InfoMax, which encourages input-to-representation information, with the Information Bottleneck, which compresses representations, assigning them to common and task-specific knowledge, respectively.Source: §2 · p. 5

Multi-task learning

  • Unlike conventional multi-task learning approaches that decompose shared and task-specific processing, this work decomposes a pretrained model into task-organized knowledge modules.Source: §2 · p. 5

3 Method

The method factorizes a multi-task teacher into task-specific students with shared and specialized modules, then uses mutual-information objectives to disentangle their representations.

3.1 Knowledge Factorization in Neural Network

  • Knowledge Factorization turns a multi-task teacher into independent single-task factor networks, preserving both task-specific knowledge and a shared common feature.Source: §3.1 · pp. 5–6

3.2 Structural Factorization

  • Each task network pairs a shared CKN for task-agnostic features with an exclusive TSN; teacher-mimicking, supervised training supports assembling networks without retraining.Source: §3.2 · p. 6

3.3 Representation Factorization

  • IMB maximizes I(Tj,Yj) and I(X,Z) while minimizing I(X,Tj), encouraging predictive task features and input-rich common features.The task-specific representation is intended to preserve only task-related information, while the common representation retains as much input information as possible.Source: §3.3 · p. 7

3.4 Variational Bound for Mutual Information

  • Because continuous-variable mutual information is difficult to estimate, the authors approximate IMB with a variational lower bound that is tight when the variational distributions match their targets.Source: §3.4 · p. 7
  • The bound uses a task-prediction head and an energy-based critic for common features, while task features use Gaussian posteriors and a zero-mean, unit-variance prior.The common-feature estimate compares teacher and CKN intermediate features using a feature-wise loss and a dimension-aligning feed-forward network.Source: §3.4 · p. 8
  • Training minimizes a combined loss for structural and representation factorization, with λI weighting the IMB objective.Source: Training · p. 8

4 Experiments

The experiments evaluate factorized networks for task prediction, disentanglement, and transfer learning across several benchmarks. They report improved task performance, more disentangled representations, and benefits from reusing common-knowledge representations.

4.1 Factor Networks Make Strong Task Prediction

  • 1.08% and 1.31%: 11-task KF improved over ResNet18-KD on ImageNet when learning from ResNet-50 and ResNet-34, respectively; KF also achieved strong student performance against other baselines on dense prediction.The dense-prediction comparison includes PASCAL, where multi-task baselines degraded relative to single-task competitors.Source: Table 2 · p. 10, §4.1 · p. 11, Table 4 · p. 12
  • In pseudo-multi-task classification, conventional KD failed to improve results, while KF's shared CKN alleviated class-imbalance effects; KF also addressed negative transfer on PASCAL through independent TSNs.The pseudo-multi-task setup assigns each network one super-class and treats other classes as negatives.Source: §4.1 · pp. 10–11

4.2 Factorization brings Disentanglement

  • KF factor networks generally scored best on disentanglement metrics, with dSprites median improvements of 0.47 on DCI and 0.09 on MIG.The evaluation compares single-task, multi-task, KD, and KF representations on dSprites and Shape3D.Source: §4.2 · p. 12, Figure 4 · p. 13
  • CKA results indicate that factorized TSNs capture purer task-specific knowledge than MTL2STL students, whose features retain higher similarity to peers.The analysis covers dSprites, Shape3D, CIFAR10, and NYUDv2.Source: §4.2 · p. 13, Figure 5 · p. 13

4.3 Common Knowledge benefits Transferring

  • ImageNet-pretrained ResNet-18 CKN weights improved performance on both CUB and Scene over original pretrained weights, and the finetuned CKN served as a better teacher for student networks.The transfer evaluation reports finetuning and distillation performance with different pretrained weights.Source: §4.3 · pp. 13–14, Table 5 · p. 14

5 Conclusion

Knowledge Factorization decomposes a pretrained teacher into task-disentangled factor networks containing task-specific and common knowledge. The paper uses an InfoMax Bottleneck objective and reports strong performance, disentanglement, modularity, and transferability across benchmarks.

5 Conclusion

  • Knowledge Factorization decomposes a pretrained teacher into task-disentangled factor networks containing task-specific and common knowledge.The factor networks may operate independently or be integrated into multi-task networks.Source: §5 · p. 14
  • The InfoMax Bottleneck objective disentangles common and task-specific representations by optimizing mutual information between input and representations.Source: §5 · p. 14
  • Across benchmarks, the method demonstrates strong and robust performance, disentanglement capability, modularity, and transferability.Source: §5 · p. 14

Related papers

Loading 2207.03337v2…