Kai Zhu 0004

dblp:75/4078-4 · DBLP profile ↗
← Back
28ranked-venue papers
6as first author
26since 2021 · last 2025
0000-0001-5412-5392ORCID · conflict

Domains — the database's venue-derived domains; a paper can count in several

Artificial intelligence and machine learning · 22 · 5 first-author · 20 since 2021Graphics, computer vision, multimedia, augmented reality and games · 18 · 6 first-author · 16 since 2021Systems, architecture and hardware · 1 · 1 since 2021Software engineering, systems software and programming languages · 1 · 1 since 2021Databases, data management, data science and information retrieval · 1 · 1 since 2021
YearPublicationVenuePosition
2025 MangaNinja: Line Art Colorization with Precise Reference Following
abstract
Derived from diffusion models, MangaNinja specializes in the task of reference-guided line art colorization. We incorporate two thoughtful designs to ensure precise character detail transcription, including a patch shuffling module to facilitate correspondence learning between the reference color image and the target line art, and a point-driven control scheme to enable fine-grained color matching. Experiments on a self-collected benchmark demonstrate the superiority of our model over current solutions in terms of precise colorization. We further showcase the potential of the proposed interactive point control in handling challenging cases (e.g., extreme poses and shadows), cross-character colorization, multi-reference harmonization, etc., beyond the reach of existing algorithms. Our code and model could be found here.
Ka Leong Cheng, Xi Chen 0042, Jie Xiao 0002, Hao Ouyang, Kai Zhu 0004, Yu Liu 0063, Yujun Shen, Qifeng Chen 0001, Ping Luo 0002
CVPR6
2025 Improved Video VAE for Latent Video Diffusion Model
abstract
Variational Autoencoder (VAE) aims to compress pixel data into low-dimensional latent space, playing an important role in OpenAI’s Sora and other latent video diffusion generation models. While most existing video VAEs inflate a pre-trained image VAE into the 3D causal structure for temporal-spatial compression, this paper presents two astonishing findings: (1) The initialization from a well-trained image VAE with the same latent dimensions is not an optimal scheme. (2) The adoption of causal reasoning leads to unequal information interactions and unbalanced performance between frames. To alleviate these problems, we propose a keyframe-based temporal compression (KTC) architecture and a group causal convolution (GCConv) module to further improve video VAE (IV-VAE). Specifically, the KTC architecture divides the latent space into two branches, in which one half completely inherits the compression prior of keyframes from a lower-dimension image VAE while the other half involves temporal compression through 3D group causal convolution, reducing temporal-spatial conflicts and accelerating the convergence speed of video VAE. The GC-Conv in the above 3D half uses standard convolution within each frame group to ensure inter-frame equivalence, and employs causal logical padding between groups to maintain flexibility in processing variable frame video. Extensive experiments on five benchmarks demonstrate the SOTA video reconstruction and generation abilities of our IV-VAE.
Pingyu Wu, Kai Zhu 0004, Yu Liu 0063, Wei Zhai, Yang Cao 0010, Zhengjun Zha
CVPR2
2025 BACON: Improving Clarity of Image Captions via Bag-of-Concept Graphs
abstract
Advancements in large Vision-Language Models have brought precise, accurate image captioning, vital for advancing multi-modal image understanding and processing. Yet these captions often carry lengthy, intertwined contexts that are difficult to parse and frequently overlook essential cues, posing a great barrier for models like GroundingDINO and SDXL, which lack the strong text encoding and syntax analysis needed to fully leverage dense captions. To address this, we propose BACON, a prompting method that breaks down VLM-generated captions into disentangled, structured elements such as objects, relationships, styles, and themes. This approach not only minimizes confusion from handling complex contexts but also allows for efficient transfer into a JSON dictionary, enabling models without linguistic processing capabilities to easily access key information. We annotated 100,000 image-caption pairs using BACON with GPT-4V and trained an LLaVA captioner on this dataset, enabling it to produce BACON-style captions without relying on costly GPT-4V. Evaluations of overall quality, precision, and recall—as well as user studies—demonstrate that the resulting caption model consistently outperforms other SOTA VLM models in generating high-quality captions. Besides, we show that BACON-style captions exhibit better clarity when applied to various models, enabling them to accomplish previously unattainable tasks or surpass existing SOTA solutions without training. For example, BACON-style captions help GroundingDINO achieve 1.51× higher recall scores on open-vocabulary object detection tasks compared to leading methods.
Zhantao Yang, Ruili Feng, Huangji Wang, Zhicai Wang, Shangwen Zhu, Han Zhang 0010, Jie Xiao 0002, Pingyu Wu, Kai Zhu 0004, Jixuan Chen, Chen-Wei Xie, Hongyang Zhang 0001, Yu Liu 0063, Fan Cheng 0002
CVPR10
2025 FastAno: Accelerating Defect Image Generation with Efficient Sampling
abstract
Defect inspection faces the challenge of insufficient data. Although existing defect generation methods can produce high-quality defect images, the time-consuming generation process hinders the online availability. To solve it, we propose FastAno, a four-step sampling model for rapid defect generation. Specifically, we first introduce the Adaptive Defect-specific Loss, which calculates region-weighted feature loss to enhance shortcut mapping of defect distribution. Secondly, we propose the Dynamic Attention Optimization Strategy, which enhances the attention activation of the anomaly semantics to improve the generation of defects, while adaptively suppressing the activation of normal semantics to mitigate the degradation of non-defect regions. Extensive experiments on MVTec AD dataset demonstrate that our method achieves significantly faster generation speed while maintaining high generation quality.
Haoyu Guan, Qianzi Yu, Kai Zhu 0004, Yang Cao 0010, Yu Kang 0001
ICME3
2025 INFP: INdustrial Video Anomaly Detection via Frequency Prioritization
abstract
Industrial video anomaly detection aims to perform real-time analysis of video streams from industrial production lines and provide anomaly alerts. Conventional video anomaly detection methods focus more on the overall image, as they aim to identify anomalies among multiple normal samples appearing simultaneously. However, industrial scenarios, where the primary focus is on a single type of product, require attention to local areas to capture fine-grained details and specific patterns. Directly applying conventional methods to industrial scenarios can result in an inability to focus on products moving along fixed trajectories, ineffective utilization of their equidistant periodicity, and greater susceptibility to lighting variations. To address these issues, we propose FreqNet, an encoder-decoder framework that learns frequency-domain features from videos to capture periodic and dynamic characteristics, enhancing the model's robustness. Specifically, a trajectory filter is proposed that takes advantage of the significant difference between moving objects and static backgrounds in the frequency domain by assigning higher weights to fixed moving trajectories. Moreover, a multi-feature fusion module is proposed, in which the frequency domain features of the video are first extracted to leverage the unique equidistant periodicity information of videos from industrial production lines. The extracted frequency domain features are subsequently fused with spatio-temporal features and contextual information is further integrated from the fused representation, effectively mitigating the impact of lighting variations on production lines. Extensive experiments on the benchmark IPAD dataset demonstrate the superiority of our proposed method over the state-of-the-art.
Qianzi Yu, Kai Zhu 0004, Yang Cao 0010, Yu Kang 0001
IJCAI2
2025 ViewPoint: Panoramic Video Generation with Pretrained Diffusion Models
abstract
Panoramic video generation aims to synthesize 360-degree immersive videos, holding significant importance in the fields of VR, world models, and spatial intelligence. Existing works fail to synthesize high-quality panoramic videos due to the inherent modality gap between panoramic data and perspective data, which constitutes the majority of the training data for modern diffusion models. In this paper, we propose a novel framework utilizing pretrained perspective video models for generating panoramic videos. Specifically, we design a novel panorama representation named ViewPoint map, which possesses global spatial continuity and fine-grained visual details simultaneously. With our proposed Pano-Perspective attention mechanism, the model benefits from pretrained perspective priors and captures the panoramic spatial correlations of the ViewPoint map effectively. Extensive experiments demonstrate that our method can synthesize highly dynamic and spatially consistent panoramic videos, achieving state-of-the-art performance and surpassing previous methods.
Zixun Fang, Kai Zhu 0004, Yu Liu 0063, Wei Zhai, Yang Cao 0010, Zhengjun Zha
NeurIPS2
2024 Hypercorrelation Evolution for Video Class-Incremental Learning
abstract
Video class-incremental learning aims to recognize new actions while restricting the catastrophic forgetting of old ones, whose representative samples can only be saved in limited memory. Semantically variable subactions are susceptible to class confusion due to data imbalance. While existing methods address the problem by estimating and distilling the spatio-temporal knowledge, we further explores that the refinement of hierarchical correlations is crucial for the alignment of spatio-temporal features. To enhance the adaptability on evolved actions, we proposes a hierarchical aggregation strategy, in which hierarchical matching matrices are combined and jointly optimized to selectively store and retrieve relevant features from previous tasks. Meanwhile, a correlation refinement mechanism is presented to reinforce the bias on informative exemplars according to online hypercorrelation distribution. Experimental results demonstrate the effectiveness of the proposed method on three standard video class-incremental learning benchmarks, outperforming state-of-the-art methods. Code is available at: https://github.com/Lsen991031/HCE
Sen Liang, Kai Zhu 0004, Wei Zhai, Yang Cao 0010
AAAI2
2024 HomoFormer: Homogenized Transformer for Image Shadow Removal
abstract
The spatial non-uniformity and diverse patterns of shadow degradation conflict with the weight sharing manner of dominant models, which may lead to an unsatisfactory compromise. To tackle with this issue, we present a novel strategy from the view of shadow transformation in this paper: directly homogenizing the spatial distribution of shadow degradation. Our key design is the random shuffle operation and its corresponding inverse operation. Specifically, random shuffle operation stochastically rearranges the pixels across spatial space and the inverse operation recovers the original order. After randomly shuffling, the shadow diffuses in the whole image and the degradation appears in a homogenized way, which can be effectively processed by the local self-attention layer. Moreover, we further devise a new feed forward network with position modeling to exploit image structural information. Based on these elements, we construct the final local window based transformer named HomoFormer for image shadow removal. Our HomoFormer can enjoy the linear complexity of local transformers while bypassing challenges of non-uniformity and diversity of shadow. Extensive experiments are conducted to verify the superiority of our HomoFormer across public datasets. Code is available at https://github.com/jiexiaou/HomoFormer.
Jie Xiao 0002, Xueyang Fu, Yurui Zhu, Dong Li 0055, Jie Huang 0017, Kai Zhu 0004, Zhengjun Zha
CVPR6
2024 DreamClean: Restoring Clean Image Using Deep Diffusion Prior
abstract
Image restoration poses a garners substantial interest due to the exponential surge in demands for recovering high-quality images from diverse mobile camera devices, adverse lighting conditions, suboptimal shooting environments, and frequent image compression for efficient transmission purposes. Yet this problem gathers significant challenges as people are blind to the type of restoration the images suffer, which, is usually the case in real-day scenarios and is most urgent to solve for this field. Current research, however, heavily relies on prior knowledge of the restoration type, either explicitly through rules or implicitly through the availability of degraded-clean image pairs to define the restoration process, and consumes considerable effort to collect image pairs of vast degradation types. This paper introduces DreamClean, a training-free method that needs no degradation prior knowledge but yields high-fidelity and generality towards various types of image degradation. DreamClean embeds the degraded image back to the latent of pre-trained diffusion models and re-sample it through a carefully designed diffusion process that mimics those generating clean images. Thanks to the rich image prior in diffusion models and our novel Variance Preservation Sampling (VPS) technique, DreamClean manages to handle various different degradation types at one time and reaches far more satisfied final quality than previous competitors. DreamClean relies on elegant theoretical supports to assure its convergence to clean image when VPS has appropriate parameters, and also enjoys superior experimental performance over various challenging tasks that could be overwhelming for previous methods when degradation prior is unavailable.
Jie Xiao 0002, Ruili Feng, Han Zhang 0010, Zhantao Yang, Yurui Zhu, Xueyang Fu, Kai Zhu 0004, Yu Liu 0063, Zhengjun Zha
ICLR8
2024 Lipschitz Singularities in Diffusion Models
abstract
Diffusion models, which employ stochastic differential equations to sample images through integrals, have emerged as a dominant class of generative models. However, the rationality of the diffusion process itself receives limited attention, leaving the question of whether the problem is well-posed and well-conditioned. In this paper, we uncover a vexing propensity of diffusion models: they frequently exhibit the infinite Lipschitz near the zero point of timesteps. We provide theoretical proofs to illustrate the presence of infinite Lipschitz constants and empirical results to confirm it. The Lipschitz singularities pose a threat to the stability and accuracy during both the training and inference processes of diffusion models. Therefore, the mitigation of Lipschitz singularities holds great potential for enhancing the performance of diffusion models. To address this challenge, we propose a novel approach, dubbed E-TSDM, which alleviates the Lipschitz singularities of the diffusion model near the zero point. Remarkably, our technique yields a substantial improvement in performance. Moreover, as a byproduct of our method, we achieve a dramatic reduction in the Fréchet Inception Distance of acceleration methods relying on network Lipschitz, including DDIM and DPM-Solver, by over 33\%. Extensive experiments on diverse datasets validate our theory and method. Our work may advance the understanding of the general diffusion process, and also provide insights for the design of diffusion models.
Zhantao Yang, Ruili Feng, Han Zhang 0010, Yujun Shen, Kai Zhu 0004, Lianghua Huang, Yu Liu 0063, Deli Zhao, Jingren Zhou 0001, Fan Cheng 0002
ICLR5
2024 CCM: Real-Time Controllable Visual Content Creation Using Text-to-Image Consistency Models
abstract
Consistency Models (CMs) have showed a promise in creating high-quality images with few steps. However, the way to add new conditional controls to the pre-trained CMs has not been explored. In this paper, we explore the pivotal subject of leveraging the generative capacity and efficiency of consistency models to facilitate controllable visual content creation via ControlNet. First, it is observed that ControlNet trained for diffusion models (DMs) can be directly applied to CMs for high-level semantic controls but sacrifice image low-level details and realism. To tackle with this issue, we develop a CMs-tailored training strategy for ControlNet using the consistency training. It is substantiated that ControlNet can be successfully established through the consistency training technique. Besides, a unified adapter can be trained utilizing the consistency training, which enhances the adaptation of DM’s ControlNet. We quantitatively and qualitatively evaluate all strategies across various conditional controls, including sketch, hed, canny, depth, human pose, low-resolution image and masked image, with the pre-trained text-to-image latent consistency models.
Jie Xiao 0002, Kai Zhu 0004, Han Zhang 0010, Yujun Shen, Zhantao Yang, Ruili Feng, Yu Liu 0063, Xueyang Fu, Zhengjun Zha
ICML2
2024 Background Activation Suppression for Weakly Supervised Object Localization and Semantic Segmentation
Wei Zhai, Pingyu Wu, Kai Zhu 0004, Yang Cao 0010, Feng Wu 0001, Zhengjun Zha
Int. J. Comput. Vis.3
2024 TF²: Few-Shot Text-Free Training-Free Defect Image Generation for Industrial Anomaly Inspection
abstract
Anomaly inspection aims at identifying various defects in real time on modern industrial production lines. However, due to insufficient anomaly data, existing detectors cannot effectively accomplish the classification of defects, thereby failing to provide guidance for subsequent production. To address it, we propose TF2, a few-shot text-free training-free defect image generation method, which jointly models the image distribution of class-agnostic defects and backgrounds, achieving efficient semantic enhancement. Firstly, we propose the Response Alignment Strategy, which merges the reversed latent space of both defect-free and defective samples, generating new defect images not limited to textual descriptions yet with consistent content. Moreover, we introduce the Defect Moving Strategy and the Regional Average Loss to merge the reversed latent space of the moving areas and enhance the variability of detail features, increasing both the location and content diversity of defects. Extensive experiments demonstrate the superiority of our model over the state-of-the-art competitors. The metrics indicate that our generated anomaly data focuses on balancing both image quality and diversity, effectively improving the performance of downstream anomaly inspection tasks.
Qianzi Yu, Kai Zhu 0004, Yang Cao 0010, Feijie Xia, Yu Kang 0001
IEEE Trans. Circuits Syst. Video Technol.2
2023 Neural Dependencies Emerging from Learning Massive Categories
abstract
This work presents two astonishing findings on neural networks learned for large-scale image classification. 1) Given a well-trained model, the logits predicted for some category can be directly obtained by linearly combining the predictions of a few other categories, which we call neural dependency. 2) Neural dependencies exist not only within a single model, but even between two independently learned models, regardless of their architectures. Towards a theoretical analysis of such phenomena, we demonstrate that identifying neural dependencies is equivalent to solving the Covariance Lasso (CovLasso) regression problem proposed in this paper. Through investigating the properties of the problem solution, we confirm that neural dependency is guaranteed by a redundant logit covariance matrix, which condition is easily met given massive categories, and that neural dependency is highly sparse, implying that one category correlates to only a few others. We further empirically show the potential of neural dependencies in understanding internal data correlations, generalizing models to unseen categories, and improving model robustness with a dependency-derived regularizer. Code to reproduce the results in this paper is available at https://github.com/RuiLiFengiNeural-Dependencies.
Ruili Feng, Kecheng Zheng, Kai Zhu 0004, Yujun Shen, Jian Zhao 0018, Deli Zhao, Jingren Zhou 0001, Michael I. Jordan, Zhengjun Zha
CVPR3
2023 Uncertainty-Aware Optimal Transport for Semantically Coherent Out-of-Distribution Detection
abstract
Semantically coherent out-of-distribution (SCOOD) detection aims to discern outliers from the intended data distribution with access to unlabeled extra set. The coexistence of in-distribution and out-of-distribution samples will exacerbate the model overfitting when no distinction is made. To address this problem, we propose a novel uncertainty-aware optimal transport scheme. Our scheme consists of an energy-based transport (ET) mechanism that estimates the fluctuating cost of uncertainty to promote the assignment of semantic-agnostic representation, and an inter-cluster extension strategy that enhances the discrimination of semantic property among different clusters by widening the corresponding margin distance. Furthermore, a T-energy score is presented to mitigate the magnitude gap between the parallel transport and classifier branches. Extensive experiments on two standard SCOOD benchmarks demonstrate the above-par OOD detection performance, outperforming the state-of-the-art methods by a margin of 27.69% and 34.4% on FPR@95, respectively. Code is available at https://github.com/LuFan3/IET-OOD.
Kai Zhu 0004, Wei Zhai, Kecheng Zheng, Yang Cao 0010
CVPR2
2023 Self-Organizing Pathway Expansion for Non-Exemplar Class-Incremental Learning
abstract
Non-exemplar class-incremental learning aims to recognize both the old and new classes without access to old class samples. The conflict between old and new class optimization is exacerbated since the shared neural pathways can only be differentiated by the incremental samples. To address this problem, we propose a novel self-organizing pathway expansion scheme. Our scheme consists of a class-specific pathway organization strategy that reduces the coupling of optimization pathway among different classes to enhance the independence of the feature representation, and a pathway-guided feature optimization mechanism to mitigate the update interference between the old and new classes. Extensive experiments on four datasets demonstrate significant performance gains, outperforming the state-of-the-art methods by a margin of 1%, 3%, 2% and 2%, respectively.
Kai Zhu 0004, Kecheng Zheng, Ruili Feng, Deli Zhao, Yang Cao 0010, Zhengjun Zha
ICCV1
2023 Regularized Mask Tuning: Uncovering Hidden Knowledge in Pre-trained Vision-Language Models
abstract
Prompt tuning and adapter tuning have shown great potential in transferring pre-trained vision-language models (VLMs) to various downstream tasks. In this work, we design a new type of tuning method, termed as regularized mask tuning, which masks the network parameters through a learnable selection. Inspired by neural pathways, we argue that the knowledge required by a downstream task already exists in the pre-trained weights but just gets concealed in the upstream pre-training stage. To bring the useful knowledge back into light, we first identify a set of parameters that are important to a given downstream task, then attach a binary mask to each parameter, and finally optimize these masks on the downstream data with the parameters frozen. When updating the mask, we introduce a novel gradient dropout strategy to regularize the parameter selection, in order to prevent the model from forgetting old knowledge and overfitting the downstream data. Experimental results on 11 datasets demonstrate the consistent superiority of our method over previous alternatives. It is noteworthy that we manage to deliver 18.73% performance improvement compared to the zero-shot CLIP via masking an average of only 2.56% parameters. Furthermore, our method is synergistic with most existing parameter-efficient tuning methods and can boost the performance on top of them. Project page can be found here.
Kecheng Zheng, Ruili Feng, Kai Zhu 0004, Jiawei Liu 0001, Deli Zhao, Zhengjun Zha, Wei Chen 0001, Yujun Shen
ICCV4
2023 Cones: Concept Neurons in Diffusion Models for Customized Generation
abstract
Human brains respond to semantic features of presented stimuli with different neurons. This raises the question of whether deep neural networks admit a similar behavior pattern. To investigate this phenomenon, this paper identifies a small cluster of neurons associated with a specific subject in a diffusion model. We call those neurons the concept neurons. They can be identified by statistics of network gradients to a stimulation connected with the given subject. The concept neurons demonstrate magnetic properties in interpreting and manipulating generation results. Shutting them can directly yield the related subject contextualized in different scenes. Concatenating multiple clusters of concept neurons can vividly generate all related concepts in a single image. Our method attains impressive performance for multi-subject customization, even four or more subjects. For large-scale applications, the concept neurons are environmentally friendly as we only need to store a sparse cluster of int index instead of dense float32 parameter values, reducing storage consumption by 90% compared with previous customized generation methods. Extensive qualitative and quantitative studies on diverse scenarios show the superiority of our method in interpreting and manipulating diffusion models.
Ruili Feng, Kai Zhu 0004, Kecheng Zheng, Yu Liu 0063, Deli Zhao, Jingren Zhou 0001, Yang Cao 0010
ICML3
2023 Customizable Image Synthesis with Multiple Subjects
abstract
Synthesizing images with user-specified subjects has received growing attention due to its practical applications. Despite the recent success in single subject customization, existing algorithms suffer from high training cost and low success rate along with increased number of subjects. Towards controllable image synthesis with multiple subjects as the constraints, this work studies how to efficiently represent a particular subject as well as how to appropriately compose different subjects. We find that the text embedding regarding the subject token already serves as a simple yet effective representation that supports arbitrary combinations without any model tuning. Through learning a residual on top of the base embedding, we manage to robustly shift the raw subject to the customized subject given various text conditions. We then propose to employ layout, a very abstract and easy-to-obtain prior, as the spatial guidance for subject arrangement. By rectifying the activations in the cross-attention map, the layout appoints and separates the location of different subjects in the image, significantly alleviating the interference across them. Using cross-attention map as the intermediary, we could strengthen the signal of target subjects and weaken the signal of irrelevant subjects within a certain region, significantly alleviating the interference across subjects. Both qualitative and quantitative experimental results demonstrate our superiority over state-of-the-art alternatives under a variety of settings for multi-subject customization.
Yujun Shen, Kecheng Zheng, Kai Zhu 0004, Ruili Feng, Yu Liu 0063, Deli Zhao, Jingren Zhou 0001, Yang Cao 0010
NeurIPS5
2023 BladeDISC: Optimizing Dynamic Shape Machine Learning Workloads via Compiler Approach
abstract
Compiler optimization plays an increasingly important role to boost the performance of machine learning models for data processing and management. With increasingly complex data, the dynamic tensor shape phenomenon emerges for ML models. However, existing ML compilers either can only handle static shape models or expose a series of performance problems for both operator fusion optimization and code generation in dynamic shape scenes. This paper tackles the main challenges of dynamic shape optimization: the fusion optimization without shape value, and code generation supporting arbitrary shapes. To tackle the fundamental challenge of the absence of shape values, it systematically abstracts and excavates the shape information and designs a cross-level symbolic shape representation. With the insight that what fusion optimization relies upon is tensor shape relationships between adjacent operators rather than exact shape values, it proposes the dynamic shape fusion approach based on shape information propagation. To generate code that adapts to arbitrary shapes efficiently, it proposes a compile-time and runtime combined code generation approach. Finally, it presents a complete optimization pipeline for dynamic shape models and implements an industrial-grade ML compiler, named BladeDISC. The extensive evaluation demonstrates that BladeDISC outperforms PyTorch, TorchScript, TVM, ONNX Runtime, XLA, Torch Inductor (dynamic shape), and TensorRT by up to 6.95×, 6.25×, 4.08×, 2.04×, 2.06×, 7.92×, and 4.16× (3.54×, 3.12×, 1.95×, 1.47×, 1.24×, 2.93×, and 1.46× on average) in terms of end-to-end inference speedup on the A10 and T4 GPU, respectively. BladeDISC's source code is publicly available at https://github.com/alibaba/BladeDISC.
Zhen Zheng, Zaifeng Pan, Dalin Wang, Kai Zhu 0004, Wenyi Zhao, Tianyou Guo, Xiafei Qiu, Minmin Sun, Feng Zhang 0007, Xiaoyong Du 0001, Jidong Zhai, Wei Lin 0016
Proc. ACM Manag. Data4
2022 AStitch: enabling a new multi-dimensional optimization space for memory-intensive ML training and inference on modern SIMT architectures
abstract
This work reveals that memory-intensive computation is a rising performance-critical factor in recent machine learning models. Due to a unique set of new challenges, existing ML optimizing compilers cannot perform efficient fusion under complex two-level dependencies combined with just-in-time demand. They face the dilemma of either performing costly fusion due to heavy redundant computation, or skipping fusion which results in massive number of kernels. Furthermore, they often suffer from low parallelism due to the lack of support for real-world production workloads with irregular tensor shapes. To address these rising challenges, we propose AStitch, a machine learning optimizing compiler that opens a new multi-dimensional optimization space for memory-intensive ML computations. It systematically abstracts four operator-stitching schemes while considering multi-dimensional optimization objectives, tackles complex computation graph dependencies with novel hierarchical data reuse, and efficiently processes various tensor shapes via adaptive thread mapping. Finally, AStitch provides just-in-time support incorporating our proposed optimizations for both ML training and inference. Although AStitch serves as a stand-alone compiler engine that is portable to any version of TensorFlow, its basic ideas can be generally applied to other ML frameworks and optimization compilers. Experimental results show that AStitch can achieve an average of 1.84x speedup (up to 2.73x) over the state-of-the-art Google's XLA solution across five production workloads. We also deploy AStitch onto a production cluster for ML workloads with thousands of GPUs. The system has been in operation for more than 10 months and saves about 20,000 GPU hours for 70,000 tasks per week.
Zhen Zheng, Xuanda Yang, Pengzhan Zhao, Guoping Long, Kai Zhu 0004, Feiwen Zhu, Wenyi Zhao, Jun Yang 0052, Jidong Zhai, Shuaiwen Song, Wei Lin 0016
ASPLOS5
2022 Self-Sustaining Representation Expansion for Non-Exemplar Class-Incremental Learning
abstract
Non-exemplar class-incremental learning is to recognize both the old and new classes when old class samples cannot be saved. It is a challenging task since representation optimization and feature retention can only be achieved under supervision from new classes. To address this problem, we propose a novel self-sustaining representation expansion scheme. Our scheme consists of a structure reorganization strategy that fuses main-branch expansion and side-branch updating to maintain the old features, and a main-branch distillation scheme to transfer the invariant knowledge. Furthermore, a prototype selection mechanism is proposed to enhance the discrimination between the old and new classes by selectively incorporating new samples into the distillation process. Extensive experiments on three benchmarks demonstrate significant incremental performance, outperforming the state-of-the-art methods by a margin of 3%, 3% and 6%, respectively.
Kai Zhu 0004, Wei Zhai, Yang Cao 0010, Jiebo Luo 0001, Zhengjun Zha
CVPR1
2022 AS-Net: Class-Aware Assistance and Suppression Network for Few-Shot Learning
Ruijing Zhao, Kai Zhu 0004, Yang Cao 0010, Zhengjun Zha
MMM (2)2
2022 Uncertainty-Aware Hierarchical Refinement for Incremental Implicitly-Refined Classification
abstract
Incremental implicitly-refined classification task aims at assigning hierarchical labels to each sample encountered at different phases. Existing methods tend to fail in generating hierarchy-invariant descriptors when the novel classes are inherited from the old ones. To address the issue, this paper, which explores the inheritance relations in the process of multi-level semantic increment, proposes an Uncertainty-Aware Hierarchical Refinement (UAHR) scheme. Specifically, our proposed scheme consists of a global representation extension strategy that enhances the discrimination of incremental representation by widening the corresponding margin distance, and a hierarchical distribution alignment strategy that refines the distillation process by explicitly determining the inheritance relationship of the incremental class. Particularly, the shifting subclasses are corrected under the guidance of hierarchical uncertainty, ensuring the consistency of the homogeneous features. Extensive experiments on widely used benchmarks (i.e., IIRC-CIFAR, IIRC-ImageNet-lite, IIRC-ImageNet-Subset, and IIRC-ImageNet-full) demonstrate the superiority of our proposed method over the state-of-the-art approaches.
Jian Yang 0003, Kai Zhu 0004, Kecheng Zheng, Yang Cao 0010
NeurIPS2
2021 Self-Promoted Prototype Refinement for Few-Shot Class-Incremental Learning
abstract
Few-shot class-incremental learning is to recognize the new classes given few samples and not forget the old classes. It is a challenging task since representation optimization and prototype reorganization can only be achieved under little supervision. To address this problem, we propose a novel incremental prototype learning scheme. Our scheme consists of a random episode selection strategy that adapts the feature representation to various generated incremental episodes to enhance the corresponding extensibility, and a self-promoted prototype refinement mechanism which strengthens the expression ability of the new classes by explicitly considering the dependencies among different classes. Particularly, a dynamic relation projection module is proposed to calculate the relation matrix in a shared embedding space and leverage it as the factor for bootstrapping the update of prototypes. Extensive experiments on three benchmark datasets demonstrate the above-par incremental performance, outperforming state-of-the-art methods by a margin of 13%, 17% and 11%, respectively.
Kai Zhu 0004, Yang Cao 0010, Wei Zhai, Zhengjun Zha
CVPR1
2021 One-Shot Texture Retrieval Using Global Grouping Metric
abstract
Texture retrieval is widely used in the fields of fashion and e-commerce. This paper presents the problem of one-shot texture retrieval: given an example of a new reference texture, we aim to detect and segment all pixels of the same texture category within an arbitrary image. To address this problem, an OS-TR network is proposed to encode both reference and query images into a texture representation space, and a better comparison is made based on the global grouping information. Because the learned texture representation should be invariant to the spatial layout while preserving the rough semantic concepts, we introduce an adaptive directionality-aware module to finely discriminate the orderless texture details. To make full use of the global context information given only a few examples, we incorporate a grouping-attention mechanism into the relation network, resulting in the per-channel modulation of the local relation features. Extensive experiments on two benchmark datasets (i.e., the DTD and ADE20K dataset) and real scenarios demonstrate that our proposed method can achieve above-par segmentation performance and robust generalization across domains.
Kai Zhu 0004, Yang Cao 0010, Wei Zhai, Zhengjun Zha
IEEE Trans. Multim.1
2020 Self-Supervised Tuning for Few-Shot Segmentation
abstract
Few-shot segmentation aims at assigning a category label to each image pixel with few annotated samples. It is a challenging task since the dense prediction can only be achieved under the guidance of latent features defined by sparse annotations. Existing meta-learning based method tends to fail in generating category-specifically discriminative descriptor when the visual features extracted from support images are marginalized in embedding space. To address this issue, this paper presents an adaptive tuning framework, in which the distribution of latent features across different episodes is dynamically adjusted based on a self-segmentation scheme, augmenting category-specific descriptors for label prediction. Specifically, a novel self-supervised inner-loop is firstly devised as the base learner to extract the underlying semantic features from the support image. Then, gradient maps are calculated by back-propagating self-supervised loss through the obtained features, and leveraged as guidance for augmenting the corresponding elements in the embedding space. Finally, with the ability to continuously learn from different episodes, an optimization-based meta-learner is adopted as outer loop of our proposed framework to gradually refine the segmentation results. Extensive experiments on benchmark PASCAL-5i and COCO-20i datasets demonstrate the superiority of our proposed method over state-of-the-art.
Kai Zhu 0004, Wei Zhai, Yang Cao 0010
IJCAI1
2019 One-Shot Texture Retrieval with Global Context Metric
abstract
In this paper, we tackle one-shot texture retrieval: given an example of a new reference texture, detect and segment all the pixels of the same texture category within an arbitrary image. To address this problem, we present an OS-TR network to encoding both reference patch and query image, leading to achieve texture segmentation towards the reference category. Unlike the existing texture encoding methods that integrate CNN with orderless pooling, we propose a directionality-aware network to capture the texture variations at each direction, resulting in spatially invariant representation. To segment new categories given only few examples, we incorporate a self-gating mechanism into relation network to exploit global context information for adjusting per-channel modulation weights of local relation features. Extensive experiments on benchmark texture datasets and real scenarios demonstrate the above-par segmentation performance and robust generalization across domains of our proposed method.
Kai Zhu 0004, Wei Zhai, Zhengjun Zha, Yang Cao 0010
IJCAI1