Source-linked AI summary

Cones: Concept Neurons in Diffusion Models for Customized Generation

Zhiheng Liu, Ruili Feng, Kai Zhu, Yifei Zhang, Kecheng Zheng, Yu Liu, Deli Zhao, Jingren Zhou, Yang Cao

arXiv:2303.05125v1cs.CV

TL;DR

The paper asks whether diffusion models contain neurons corresponding to particular subjects, as human brains contain concept-sensitive neurons. It identifies such neurons with gradient-based analysis and manipulates them for customized generation, including multi-subject images. The method generates up to four diverse subjects in one image and uses sparse neuron indices for storage-efficient customization.

  • Problem

    The paper investigates whether diffusion models encode different subjects separately in small neuron clusters, extending the question of human concept neurons to artificial generative models.

  • Method

    Cones identifies subject-specific parameters in diffusion-model K-V attention layers using a gradient-based criterion and generates subjects by shutting the corresponding concept neurons.

  • Results

    Further fine-tuning after concatenating concept neurons enables generation of up to four different diverse subjects in one image.

  • Takeaways & Limitations

    Concept neurons provide an interpretable and additive mechanism for customized generation, allowing clusters for multiple subjects to be concatenated directly.

  • Takeaways & Limitations

    The method limits concept-neuron search to parameters in the K-V attention layers of the diffusion model.

Abstract

from arXiv · show

Human brains respond to semantic features of presented stimuli with different neurons. It is then curious whether modern deep neural networks admit a similar behavior pattern. Specifically, this paper finds a small cluster of neurons in a diffusion model corresponding to a particular subject. 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. A few steps of further fine-tuning can enhance the multi-concept capability, which may be the first to manage to generate up to four different subjects in a single image. 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 values of the parameters, which reduces storage consumption by 90\% compared with previous subject-driven generation methods. Extensive qualitative and quantitative studies on diverse scenarios show the superiority of our method in interpreting and manipulating diffusion models.

1. Introduction

The paper asks whether diffusion models contain subject-specific concept neurons and proposes Cones, a gradient-based method to identify and manipulate them. These neurons support robust subject generation, additive multi-concept composition, and storage-efficient customization.

  • The paper investigates whether diffusion-model neurons separately encode particular subjects, analogous to concept neurons in the human brain.
  • Cones identifies subject-specific concept neurons using gradient statistics after only a few gradient computations.The method targets parameters in the attention layers of a pretrained text-to-image diffusion model.
  • Concept neurons remain effective across float32, float16, quaternary, and binary digital accuracy settings.Binary accuracy requires no further training and uses minimum storage, so it becomes the default subject-driven generation method.
  • Concatenating concept-neuron clusters can generate multiple subjects together, while further fine-tuning enables up to four diverse subjects in one image.
  • Cones supports large-scale use through sparse concept-neuron storage rather than dense parameter storage.

2. Preliminaries and Background

The paper situates Cones within diffusion-based conditional generation and customized generation. It describes text-to-image diffusion models, prior-preserving subject implantation, and the objective used to bind a subject to a text identifier.

  • Diffusion Models: Diffusion models learn image distributions through gradual denoising, with conditional methods guiding generation using classifiers or paired conditioning data.Classifier-free approaches can produce detailed results but require substantial data and computation.
  • Text-to-Image Diffusion Model: Text-to-image diffusion models guide denoising with text prompts and can generate images by denoising standard Gaussian noise.The paper uses Stable Diffusion V1.4 as its default model because of its performance and availability.
  • Customized Generation: Customized generation implants a subject into a diffusion model and binds it to a unique text identifier for generating varied renditions.The subject is typically represented using three to five images taken from different viewpoints and conditions.
  • Customized Generation: Subject implantation combines a concept-preserving loss with a prior-preserving loss to retain the subject while avoiding over-fitting and language drift.
  • Customized Generation: The prior-preserving loss uses image-text pairs showing different subjects from the same category as the subject being implanted.

3. Method

Cones identifies sparse concept neurons in diffusion-model K-V attention layers using gradient-based statistics, then manipulates them to interpret and control subject generation. Shutting, concatenating, and lightly fine-tuning these neurons supports single- and multi-subject customization.

  • Concept-neuron identification: Cones searches K-V attention parameters for a sparse neuron subset whose changes alone can generate a subject across text-prompted contexts.The target subset is much smaller than the full parameter vector.
  • Concept-neuron identification: Gradient statistics identify concept neurons by testing whether scaling down a parameter decreases the concept-implanting loss.The method uses adaptive sampling and parallel mask computation, requiring only a few gradient computations.
  • Interpretability and control: Shutting subject-specific concept neurons draws the subject outline in the corresponding text-identifier attention map and generates that subject in the output.This links the neurons to both network representations and generated subjects.
  • Multi-subject generation: Concatenating concept-neuron clusters supports multi-subject generation, while fine-tuning their combined loss can generate up to four diverse subjects in one image.The concatenated representation provides a practical starting point for resolving subtle conflicts between independently computed clusters.

A V1* woman wearing V2*

Cones represents customized subjects with sparse concept-neuron indexes rather than full diffusion-model parameters, reducing storage requirements for large-scale use.

  • Cones stores sparse concept-neuron indexes instead of full-precision diffusion-model parameters.The indexes use int data types rather than float representations.
  • Cones requires no more than 10% of the memory used by previous customized-generation methods.

4. Experiments

Experiments evaluate Cones across qualitative, quantitative, user-study, tuning-free, and storage settings. Results show stronger multi-subject generation, competitive single-subject performance, and substantially lower storage costs.

  • Qualitative Evaluation: Figure 7 compares Cones with Dreambooth and Custom Diffusion across single- and multi-subject prompts, including sweater, playing, sitting, wearing, and chair semantics.
  • Qualitative Evaluation: Cones generates up to four subjects in one scene while preserving good visual accuracy for all subjects.The evaluated settings include one, two, three, and four target subjects.
  • Tuning-Free Comparison: Cones outperforms Custom Diffusion's tuning-free composition in visual quality and subject-generation accuracy without further fine-tuning.Cones directly concatenates concept-neuron clusters for multiple subjects.
  • Quantitative Evaluation: As the number of subjects increases, Cones outperforms competitors on all reported metrics.For single-subject generation, Cones has higher text alignment but slightly lower visual accuracy than Dreambooth.
  • User Study: Cones receives the most user-study votes in all cases except single-subject image alignment.
  • Sparsity and Storage: Cones requires less storage than competing customized-generation methods by recording sparse integer indexes for concept neurons.The storage gap grows as more target subjects require more indexes.

5. Conclusion

Cones identifies small clusters of concept neurons in diffusion models that control particular subjects. Shutting or combining these clusters supports contextualized and multi-subject generation, with fine-tuning enabling up to four subjects in one image.

  • A small cluster of concept neurons dominates generation of a given subject in diffusion models.
  • Shutting concept neurons yields renditions of the corresponding subject in different text-prompted contexts.
  • Concatenating clusters for different subjects can generate all those subjects in one image.
  • Further fine-tuning enables generation of up to four different subjects in one image.
  • Comparisons report superiority in visual quality, semantic alignment, multi-subject generation capability, and storage consumption.

A.1. Proof to Theorem 3.1

The proof derives an accelerated procedure for identifying concept neurons by optimizing an auxiliary variable that masks the model parameters. It then converts the variable updates into a binary concept-neuron mask.

  • The proof reparameterizes the model with γ = ξ ⊙ θ and performs gradient descent on Lcon(ξ ⊙ θ) using ξ as the optimization variable.The update is initialized with ξ = 1 and uses learning rate β or ρ, depending on the formulation.
  • When the learning rate is small and the iteration count is limited, the reparameterized variable remains approximately equal to θ.
  • The procedure accelerates concept-neuron computation by avoiding direct optimization over the full parameter vector θ.
  • The concept-neuron mask is computed from the parameter-change statistic Mp and threshold τ as M = 1 − (Mp > τ) = 1 − (ξK < 1 − βτ).
  • Algorithm A2 outputs a binary mask in which 1 denotes a non-concept neuron and 0 denotes a concept neuron.

B. Experiment Setups

The experiments use Stable Diffusion V1.4 with standardized sampling and report implementation settings for the compared methods and concept-neuron extraction. The evaluation covers subject combinations from one to four subjects and includes user preferences for multi-subject generation.

  • Experimental settings: All methods use Stable Diffusion V1.4, 50 DPM-Solver steps, and scale 7.5 for a fair comparison.Experiments are conducted on an A-100 GPU.
  • Concept-neuron extraction: Algorithm A2 identifies concept neurons by K gradient-descent updates on Lcon(ξ ⊙ θ), followed by thresholding Mp to produce a binary mask.The mask marks non-concept neurons with 1 and concept neurons with 0.
  • Evaluation protocol: User studies measure the percentage of users selecting each method's image as best, including image and text alignment for multi-subject generation.
  • Implementation settings: The compared implementations use distinct training configurations, including batch sizes, learning rates, and training-step counts for DreamBooth, Custom Diffusion, and Cones.Reported settings include 800 steps for DreamBooth, 600 for Custom Diffusion, and 1,000 steps for single-subject Cones training.
  • Evaluation protocol: The evaluation constructs 12 subject combinations across one- to four-subject generation tasks, with four prompts and 50 random seeds per combination.The best two outputs represent each prompt, and the procedure is applied to all four methods.

C.1. Sequential Training Comparison.

Cones performs better than comparison methods in sequential subject training, retaining the first subject while learning a second. Additional experiments show style expression, style fine-tuning, and improved text-guided image editing behavior.

  • Sequential training: In sequential training, Cones retains the first subject better than Custom Diffusion and DreamBooth while optimizing a second subject.The comparison evaluates the second subject while shutting the first subject's corresponding concept neurons.
  • Style control: Cones can express a specified style through text guidance and can be fine-tuned on a fixed style.
  • Text-guided editing: Cones produces generations more similar to textual descriptions when editing expressions, object colors, and backgrounds.
  • Fine-tuning effects: Fine-tuning shifts generated images toward target images and reduces diversity relative to the pretrained model, a behavior noted for the compared setting.

C.5. More results on multi subjects

Cones is evaluated on multi-subject generation, style conversion, editing, and comparisons with Textual Inversion and Custom Diffusion. The results include improved multi-subject composition when Cones incorporates Attend-and-Excite.

  • Multi-subject generation: Cones with Attend-and-Excite generates better multi-subject results than without this addition.The paper introduces Attend-and-Excite to address difficulty generating multiple subjects in one prompt.
  • Style conversion and editing: Additional results cover style conversion and image editing alongside generated subjects.The appendix figures label style conversion results and editing results separately.
  • Multi-subject generation: Textual Inversion struggles with composing multiple subjects in a single image.This comparison is shown in the multi-subject generation results.
  • Multi-subject generation: Cones generations after fine-tuning show more diversity similar to the pretrained model, whereas Custom Diffusion retains target-image characteristics without the text identifier.The comparison concerns overfitting on the training prompt template.
  • Multi-subject generation: The qualitative examples include combinations of flowers, cats, sunglasses, barrels, and paintings in multi-subject scenes.The displayed prompts include flowers next to cats, cats wearing sunglasses, and barrels decorated with flowers.
Loading 2303.05125v1…