# All-optical ultrafast ReLU function for energy-efficient nanophotonic deep
learning

Gordon H.Y. Li  Department of Applied Physics, California Institute of
Technology, Pasadena, CA 91125, USA  Ryoto Sekine  Department of Electrical
Engineering, California Institute of Technology, Pasadena, CA 91125, USA
Rajveer Nehra  Department of Electrical Engineering, California Institute of
Technology, Pasadena, CA 91125, USA  Robert M. Gray  Department of Electrical
Engineering, California Institute of Technology, Pasadena, CA 91125, USA  Luis
Ledezma  Department of Electrical Engineering, California Institute of
Technology, Pasadena, CA 91125, USA  Jet Propulsion Laboratory, California
Institute of Technology, Pasadena, California 91125, USA  Qiushi Guo
Department of Electrical Engineering, California Institute of Technology,
Pasadena, CA 91125, USA  Alireza Marandi  Department of Applied Physics,
California Institute of Technology, Pasadena, CA 91125, USA  Department of
Electrical Engineering, California Institute of Technology, Pasadena, CA
91125, USA<EMAIL_ADDRESS>
###### Abstract

In recent years, the computational demands of deep learning applications have
necessitated the introduction of energy-efficient hardware accelerators.
Optical neural networks are a promising option; however, thus far they have
been largely limited by the lack of energy-efficient nonlinear optical
functions. Here, we experimentally demonstrate an all-optical Rectified Linear
Unit (ReLU), which is the most widely used nonlinear activation function for
deep learning, using a periodically-poled thin-film lithium niobate
nanophotonic waveguide and achieve ultra-low energies in the regime of
femtojoules per activation with near-instantaneous operation. Our results
provide a clear and practical path towards truly all-optical, energy-efficient
nanophotonic deep learning.

##  1 Introduction

Over the past decade, deep learning has revolutionized many important
applications including computer vision, speech recognition, and natural
language processing [1]. However, the explosive growth of modern deep learning
models has quickly outpaced improvements in conventional von Neumann computing
architectures and ushered in the use of dedicated hardware accelerators. The
quest for ever-faster and more energy-efficient hardware for deep learning
began with exploiting the graphics processing unit (GPU), then application-
specific integrated circuits such as Google’s tensor processing unit (TPU),
and more recently the development of non-von Neumann analog architectures [2,
3]. Naturally, photonics has attracted attention as a promising candidate due
to its potential for massive parallelism and ultrafast operation [4]. Indeed,
optical neural networks (ONNs) have been experimentally demonstrated in a
variety of platforms including free-space optics [5, 6, 7, 8, 9, 10, 11],
optical fiber [12, 13, 14, 15, 16, 17], and photonic integrated circuits [18,
19, 20, 21, 22].

In general, deep neural networks require two major types of computations: (1)
linear operations in the form of matrix multiplications and convolutions,
which represent the synaptic connections of the network, and (2) nonlinear
activation functions, which represent the neuron activations. ONNs excel at
performing energy-efficient linear operations in the optical domain, which
forms the bulk of computations for deep learning. However, a major remaining
roadblock is achieving scalable energy-efficient nonlinear activation
functions, which comprises a smaller but essential part of the deep learning
workload. Thus, the majority of ONN implementations still opt to utilize
digital electronics to perform the nonlinear activation functions. In doing
so, the optoelectronic and analog-to-digital conversion typically imposes
significant speed and energy limitations. On the other hand, the demonstrated
all-optical approaches based on various processes [17, 19, 7, 13, 23] are
still too energy-intensive and/or slow compared to electronics. This is
because photon-photon interactions are typically weak and require either high
light intensities or high-Q resonant cavities, both of which are undesirable
for scalable computing purposes. An all-optical, ultrafast, and energy-
efficient nonlinear activation function is yet to be demonstrated to unlock
the full capabilities of ONNs. Such a function should also be compact, highly
scalable, and compatible with existing deep learning models.

In this work, we propose and experimentally demonstrate the first photonic
device, to the best of our knowledge, that satisfies all the aforementioned
criteria for an all-optical nonlinear activation function. It implements the
Rectified Linear Unit (ReLU) function, defined as
$\mathrm{ReLU}(x)=\mathrm{max}(0,x)$, which is one of the most widely used
nonlinear activation functions for deep learning. The widespread adoption of
the ReLU function was essential in sparking the deep learning revolution due
to its favorable properties for backpropagation training and simple
implementation in digital electronics [1]. However, its optical implementation
has remained challenging and posed a major hurdle for the real-world
applicability of ONNs.

##  2 Method

###  2.1 Principle of operation

The operating principle of our device is illustrated in Fig. 1.

Figure 1: Operating principle of the all-optical ReLU function using a
nonlinear photonic waveguide. (a) For positive inputs with phase of
$\phi_{2\omega}=+\pi/2$, the phase relationship between the signal and bias is
$2\phi_{\omega}-\phi_{2\omega}=\pi/2$, which causes SHG that depletes $\omega$
and amplifies $2\omega$. (b) For negative inputs, $\phi_{2\omega}=-\pi/2$, the
phase relationship $2\phi_{\omega}-\phi_{2\omega}=3\pi/2\rightarrow-\pi/2$
causes DOPA that amplifies $\omega$ and depletes $2\omega$.

We encode the signal information into the coherent optical field of pulses
centered at frequency $2\omega$, with positive values represented by
$\phi_{2\omega}=+\pi/2$ phase states, and negative values represented by
$\phi_{2\omega}=-\pi/2$ phase states. By co-propagating the signal pulses with
bias pulses centered at frequency $\omega$, with fixed input power and phase
at $\phi_{\omega}=+\pi/2$, we can induce different nonlinear optical effects
for the two possible $\phi_{2\omega}$ signal phases depending on the value of
the phase relationship $2\phi_{\omega}-\phi_{2\omega}$. For the positive
signal values with phase $\phi_{2\omega}=+\pi/2$, the phase relationship
yields $2\phi_{\omega}-\phi_{2\omega}=+\pi/2$. This induces second harmonic
generation (SHG), which is a $\chi^{(2)}$ nonlinear optical process that
converts two photons of frequency $\omega$ into a photon of frequency
$2\omega$, hence depleting $\omega$ and amplifying $2\omega$. Conversely, for
the negative signal values with phase $\phi_{2\omega}=-\pi/2$, the phase
relationship yields $2\phi_{\omega}-\phi_{2\omega}=3\pi/2\rightarrow-\pi/2$.
This induces degenerate optical parametric amplification (DOPA), which is the
inverse process of SHG that converts a photon of frequency $2\omega$ into two
photons of frequency $\omega$, hence depleting $2\omega$ and amplifying
$\omega$. By judiciously choosing the length and bias power, we can achieve
the desired shape of the ReLU function. We emphasize that our approach
utilizes coherent parametric processes which allows us to implement both
positive and negative values, unlike previous optical and optoelectronic
methods based on incoherent absorption processes that can only implement
positive values.

###  2.2 Device design

To implement the $\chi^{(2)}$-based ReLU function, we use a periodically poled
thin-film lithium niobate (PPLN) nanophotonic waveguide that exploits the
strong and instantaneous $\chi^{(2)}$ optical nonlinearity of lithium niobate
and tight spatial confinement of the waveguide modes to enhance the
nonlinearity [24]. Additionally, careful qausi-phase matching and dispersion
engineering enables ultra-broadband and low-energy interactions over mm-long
propagation lengths, further enhancing the nonlinear optical processes using
femtosecond laser pulses [25, 26, 27]. Images of the device are shown in Fig.
2. The PPLN nanophotonic waveguide is $L=2.5~{}\mathrm{mm}$ long and was
fabricated on a $700$-nm thick X-cut MgO-doped lithium niobate thin-film on
2-$\mathrm{\mu}$m thick SiO2 with lithium niobate substrate by dry etching
with $\mathrm{Ar}^{+}$ plasma, achieving smooth ridge side-walls with slant
angle of $\theta\approx 60^{\circ}$ as shown in Fig. 2(a). The waveguide was
electrically poled with a period of 5.17$~{}\mathrm{\mu}$m, as shown in Fig.
2(b), to ensure efficient SHG and DOPA. Dispersion engineering of the
fundamental TE mode of the ridge waveguide, shown in Fig. 2(c), allows for
negligible group velocity mismatch and group velocity dispersion of $\omega$
and $2\omega$ pulses centered at $1045~{}\mathrm{nm}$ and
$2090~{}\mathrm{nm}$, respectively. This enforces good temporal overlap of the
pulses over the entire PPLN propagation length. The ideal parameters found
from simulation were a ridge top width of $w=1700~{}\mathrm{nm}$ and etch-
depth of $h=350~{}\mathrm{nm}$. See [27] for further details about fabrication
and dispersion engineering of PPLN nanophotonic waveguides.

Figure 2: Images of the PPLN nanophotonic waveguide. (a) Scanning electron
microscope image of the ridge waveguide. (b) Two-photon absorption microscope
image of the PPLN ferroelectric domains with poling period of
5$~{}\mathrm{\mu}$m. (c) Simulated electric field distributions of the
fundamental TE modes at $1045~{}\mathrm{nm}$ ($2\omega$) and
$2090~{}\mathrm{nm}$ ($\omega$).

##  3 Results

###  3.1 Femtojoule ReLU function

The measured response of the all-optical ReLU is shown in Fig. 3. The
nonlinear function given by the PPLN was measured using a free-space chip
characterization setup. The source at $1045~{}\mathrm{nm}$ (signal) was a
Yb:fiber mode-locked laser producing $75$-fs long pulses at a $250$-MHz
repetition rate (Menlo Systems Orange). The same laser pumped a homemade
degenerate optical parametric oscillator to generate the pulses at
$2090~{}\mathrm{nm}$ (bias). The $2\omega$ and $\omega$ pulses were coupled
into and out of the PPLN using reflective objectives focused on the waveguide
facets. Finally, the relative phase of the $2\omega$ signal and $\omega$ bias
was set using a delay arm, and the power varied using a tunable attenuator.
See Supplement 1 for further details about the experimental setup.

Figure 3: Output signal pulse energy versus input signal pulse energy for both
negative and positive inputs. There is good agreement between the ideal ReLU
function (dashed black line), simulation (dashed blue line) and experimental
results (red circles) for a bias pulse energy of
$E_{\omega}(0)=270~{}\mathrm{fJ}$, and signal pulse energies of femtojoules
per activation.

Our experimental results show good agreement with the ideal ReLU function
($R^{2}=0.9895$), and demonstrates energy-efficient signal pulse energies in
the regime of femtojoules per activation. Note that the important feature of
the function is its nonlinear shape, since scaling/shifting the
horizontal/vertical directions can be accomplished with linear optical
transformations. In theory, the ideal ReLU function requires an arbitrarily
long PPLN and low bias pulse energy. However, in practice we must choose the
bias pulse energy so as to best approximate the ReLU function given our fixed
device length. Thus, there are small discrepancies around $E_{2\omega}(0)=0$,
since neither the SHG nor DOPA processes sufficiently saturate at the ultra-
low energies. The maximum cutoff pulse energy is determined by the onset of
supercontinuum generation from strong back-conversion processes, which
undesirably degrades the pulse shape. To verify that the expected device
response matches our physical picture of the operating principle, we also
performed nonlinear pulse propagation simulations of the PPLN nanophotonic
waveguide. See Supplement 1 for more details about the simulation methods.

Remarkably, we show that the PPLN nanophotonic waveguide can also approximate
other commonly used variants of the ReLU function, simply by tuning the bias
pulse energy.

Figure 4: Other variants of the ReLU function can be approximated by tuning
the bias pulse energy. For example, the (a) ELU function using bias pulse
energy of $E_{\omega}(0)=450~{}\mathrm{fJ}$ and (b) GELU function using bias
pulse energy of $E_{\omega}(0)=910~{}\mathrm{fJ}$. Ideal function curves are
shown by the dashed black lines, and experimental results with red circles.

For example, the Exponential Linear Unit (ELU) defined as $\mathrm{ELU}(x)=x$
if $x>0$ and $\mathrm{ELU}(x)=\exp{(x)}-1$ if $x<0$, which has been shown to
outperform the ReLU function in certain cases [28], is achieved using a bias
pulse energy of $E_{\omega}(0)=450~{}\mathrm{fJ}$ as shown in Fig. 4(a). In
addition, we also implement the Gaussian Error Linear Unit (GELU) defined as
$\mathrm{GELU}(x)=x\Phi(x)$ where $\Phi(x)$ is the Gaussian cumulative
distribution using a bias pulse energy of $E_{\omega}(0)=910~{}\mathrm{fJ}$ as
shown in Fig. 4(b). The GELU function is used extensively in Transformer
networks for natural language processing, which are regularly amongst the
largest deep learning models [29]. Thus, our all-optical PPLN nanophotonic
waveguide implementation gains greater real-world applicability by being
compatible with a wide range of existing deep learning models, especially the
largest models where energy efficiency is paramount. Indeed, compatibility has
been problematic in previous implementations of optical and optoelectronic
nonlinear activation functions. By alleviating this problem, we expand the
potential functionality of ONNs by avoiding the need to train new specialized
models.

###  3.2 Ultrafast time response

Ideally, the time per activation should be near-instantaneous due to the
ultrafast $\chi^{(2)}$ nonlinearity in lithium niobate. However, in practice,
the response time is limited by the finite phase-matching bandwidth as well as
non-zero group velocity mismatch, group velocity dispersion, and higher-order
dispersion terms. To determine the response time of the device, we used the
pump-probe technique commonly used to characterize all-optical switches [26,
30, 31] (see Supplement 1 for more details). In this case, the pump pulse is
the $\omega$ pulse and the probe pulse is the $2\omega$ pulse. We measured the
ultrafast ReLU dynamics by varying the time delay between the $\omega$ and
$2\omega$ pulses at a fixed pulse energy.

Figure 5: Pump-probe ultrafast timing measurements of the ReLU dynamics. The
autocorrelation (yellow circles shifted vertically for clarity) of the input
$\omega$ pulse is well-explained by a Gaussian profile (purple line) with FWHM
of $(56.4\pm 1.5)~{}\mathrm{fs}$. The pump-probe signal obtained at a fixed
pulse energy (blue circles) is fit (orange line) by convolving the input
autocorrelation with exponential growth and decay for positive and negative
time delays, respectively. The best fit yields a rise time of $(18.9\pm
1.9)~{}\mathrm{fs}$ and a fall time of $(28.4\pm 1.1)~{}\mathrm{fs}$.

Fig. 5 shows the intensity envelope of the pump-probe signal as the time delay
is varied as well as the autocorrelation of the input $\omega$ pulse. The
input autocorrelation is well-explained by a Gaussian profile with FWHM of
$(56.4\pm 1.5)~{}\mathrm{fs}$. We extract the characteristic rise and fall
times by fitting the pump-probe signal with exponential growth and decay
functions for positive and negative time delays, respectively, convolved with
the input autocorrelation. The best fit yields a rise time of $(18.9\pm
1.9)~{}\mathrm{fs}$ and a fall time of $(28.4\pm 1.1)~{}\mathrm{fs}$. This
implies that the characteristic response time of the ReLU dynamics is
$(47.3\pm 3.0)~{}\mathrm{fs}$, thus verifying that the ultrafast optical
nonlinearity is responsible for the ReLU response, and ruling out the
possibility of any slower optical nonlinearities such as photorefractive or
thermo-optic effects. Therefore, we can reasonably regard the $2\omega$ signal
pulse length of $\sim 75~{}\mathrm{fs}$ as the time per activation for the
all-optical ReLU. We note that that better dispersion engineering can lead to
even faster activation times.

###  3.3 Simulated deep learning performance

One distinct advantage of our approach is that, unlike previous all-optical
and optoelectronic nonlinear activation functions, it can faithfully reproduce
the ideal ReLU function. Therefore, we can leverage the large number of
existing pretrained deep learning models that use the ReLU function (or its
variants) for nonlinear activations. Although ONNs have been demonstrated that
accurately reproduce linear operations such as matrix multiplication and
convolution, the use of atypical nonlinear activation functions in the optical
domain has required the training of new custom deep learning models. To
improve upon this, we simulated the performance of the all-optical ReLU
function when used as part of a pretrained convolutional neural network (CNN)
for the prototypical task of MNIST handwritten digits image classification
[32].

Figure 6: Simulated deep learning performance of the experimentally measured
all-optical ReLU function for MNIST handwritten digits image classification.
(a) A pretrained CNN was used where the ideal ReLU layers are replaced with
custom layers representing the experimentally measured ReLU response (after
shifting/scaling) then fine-tuned by training for 2 epochs (batch size of 128)
to improve the test accuracy (blue line) back to the ideal pretrained model
accuracy (dashed black line). (b) Confusion matrix on the MNIST task for the
final network, which achieved $99.13\%$ test accuracy.

The MNIST dataset contains $28\times 28$ pixels gray-scale images of
handwritten digits with 50000 training samples and 10000 test samples. We used
a standard CNN architecture (see Supplement 1 for full details) containing
convolutional layers and ideal ReLU layers followed by a fully-connected layer
and softmax classification output. The pretrained CNN achieved an ideal test
accuracy of $99.13\%$. Next, the ideal ReLU layers were replaced with custom
layers representing the experimentally measured ReLU response (after proper
shifting/scaling) without changing any of the other layers. This caused a
slight drop in test accuracy to $98.8\%$ due to the slight deviations between
the experimentally measured and ideal ReLU functions. To remedy this, the CNN
was then fine-tuned by training for only 2 epochs (the CNN sees each sample
once per epoch) to regain the ideal pretrained model accuracy of $99.13\%$ as
shown in Fig. 6. Fine-tuning is necessary for any analog hardware
implementation due to unavoidable fabrication errors, noise and other non-
idealities encountered [33]. Note that this method requires far less time
compared to previous proposals for training new custom ONN models, which
required $>25$ training epochs [34, 35]. Therefore, our all-optical ReLU
provides the missing link to allow ONNs to take advantage of existing
pretrained models. We note that the softmax classification layer is yet to be
faithfully implemented in an ONN which accounts for a small portion of the
computation compared to the convolutions, matrix multiplications and ReLU
nonlinear activations.

##  4 Discussion

###  4.1 Comparison of energy and time per activation

In this section, we compare the PPLN nanophotonic waveguide to other optical
[17, 19, 13, 23], optoelectronic [11, 14, 15, 16, 21, 36, 37, 38, 39], analog
electronic [40, 41, 42], and digital electronic [43] nonlinear activation
functions to demonstrate the state-of-the-art performance of our device. In
this case, the appropriate figure of merit is the energy-time product, which
properly accounts for both the energy consumed and time taken per activation.
To quantify the energy per activation, we follow the convention in [35], as
being the energy needed to generate a $50\%$ change in the power transmission
with respect to the transmission with null input. In this case, our device has
an energy per activation of $\sim 16~{}\mathrm{fJ}$. The bias pulse energy is
not included since it is not dissipated and can be reused. The time per
activation is given by the signal pulse width of $\sim 75~{}\mathrm{fs}$,
owing to the near-instantaneous $\chi^{(2)}$ nonlinearity of lithium niobate
as explained in Section 3.2. Therefore, we achieve an energy-time product of
$1.2\times 10^{-27}~{}\mathrm{J\cdot s}$. The energy and time per activation
of our device is compared to other experimental demonstrations in Fig. 7.

Figure 7: Comparison of energy and time per activation of this work (red star)
to previous all-optical (red circle), optoelectronic (blue square), analog
electronic (green diamond), and digital electronic (magenta triangle)
nonlinear activation functions. The numeric labels show reference numbers and
dashed black lines show the energy-time product contours.

We attempted to consider device-level metrics wherever possible to provide a
fair comparison, however, we acknowledge that this was not always possible for
nonlinear activations as part of complete networks since fan-out and
cascadability constraints impose additional energy and time costs. Despite
this, the outstanding metrics of our device represents a significant
breakthrough for optical nonlinear activation functions. For state-of-the-art
digital electronics, such as the NVIDIA A100 Tensor Core GPU [44] based on
7-nm process node [45], we generously assume that the ReLU function consumes
$\sim 1~{}\mathrm{fJ}$ per activation, and occurs in a single
$1~{}\mathrm{GHz}$ clock cycle. We see that, although our device still has an
order of magnitude greater energy per activation, the time per activation is
four orders of magnitude faster. Hence, we achieve an energy-time product that
is three orders of magnitude better than state-of-the-art digital electronics.
Our numerical simulations (see Supplement 1) predict that the PPLN
nanophotonic waveguide can realistically achieve a ReLU-like response with
sub-femtojoule energy per activation. This would even surpass the energy
efficiency of state-of-the-art digital electronics. We attribute the
discrepancy between our experimental results and the theoretically predicted
limits for the energy scale to the imperfect phase-matching and fabrication
error of our device.

###  4.2 Potential network architectures

So far, we have demonstrated how PPLN nanophotonic waveguides can implement
all-optical, ultrafast, energy-efficient nonlinear activation functions, which
forms only one building block of a full neuron. In this section, we briefly
discuss how our device can be integrated into a complete ONN architecture. In
principle, the all-optical ReLU is compatible with most existing ONN
architectures that can accurately implement linear operations such as matrix
multiplication and convolutions. However, in practice, the speed bottleneck
will likely be the encoding of information into the required coherent optical
amplitudes. In this case, PPLN nanophotonic waveguides can be monolithically
integrated with high-speed electro-optic modulators in thin-film lithium
niobate, demonstrated to achieve bandwidths beyond $100~{}\mathrm{GHz}$ [46].
Furthermore, the light sources can also be integrated on-chip using thin-film
lithium niobate optical parametric oscillators [47]. Therefore, all the
fundamental building blocks needed for a complete ONN in thin-film lithium
niobate already exist. Given the rapid increases in scalability of thin-film
lithium niobate photonics, we are confident that a complete ONN can be
demonstrated in the near-future. One potential approach is to use Mach-Zehnder
interferometer meshes [18] or photonic tensor cores with waveguide cross-bar
arrays [20] to implement the linear matrix multiplications, then cascaded into
PPLN nanophotonic waveguides to perform nonlinear activations. Another
promising method is to use a time-multiplexed architecture similar to ones
demonstrated for coherent Ising machines [48] or photonic reservoir computers
[15, 14].

##  5 Conclusion

In conclusion, we have demonstrated an all-optical ultrafast ReLU function
using a PPLN nanophotonic waveguide. It has an energy per activation of $\sim
16~{}\mathrm{fJ}$ and time per activation of $\sim 75~{}\mathrm{fs}$, thus
achieving a state-of-the-art energy-time product of $1.2\times
10^{-27}~{}\mathrm{J\cdot s}$. Furthermore, we demonstrated how the same
device can be used to implement other common variants of the ReLU function,
and showed how it can exploit existing pretrained deep learning models to
greatly reduce training time. Given the simplicity of our device, and the
rapid improvements in scalability of thin-film lithium niobate photonics, we
envisage that it will be able to replace periphery digital electronic circuits
for calculating nonlinear activations in ONNs. Therefore, we have presented a
clear and practical path towards truly all-optical, energy-efficient photonic
deep learning.

## Funding

The authors gratefully acknowledge support from ARO grant no.
W911NF-18-1-0285, NSF grant no. 1846273 and 1918549, AFOSR award FA9550-
20-1-0040, and NASA/JPL. The authors wish to thank NTT Research for their
financial and technical support.

## Acknowledgments

The device nanofabrication was performed at the Kavli Nanoscience Institute
(KNI) at Caltech.

## Disclosures

GHYL, RS, RN, RMG, AM: California Institute of Technology (P)

## Data availability

Data underlying the results presented in this paper are not publicly available
at this time but may be obtained from the authors upon reasonable request.

## Supplemental document

See Supplement 1 for supporting content.

## References

  * [1] I. Goodfellow, Y. Bengio, and A. Courville, Deep learning.  MIT press, 2016.
  * [2] V. Sze, Y.-H. Chen, T.-J. Yang, and J. S. Emer, “Efficient processing of deep neural networks: A tutorial and survey,” Proceedings of the IEEE, vol. 105, no. 12, pp. 2295–2329, 2017.
  * [3] Y. LeCun, “Deep learning hardware: past, present, and future,” in 2019 IEEE International Solid-State Circuits Conference-(ISSCC), pp. 12–19, IEEE, 2019.
  * [4] G. Wetzstein, A. Ozcan, S. Gigan, S. Fan, D. Englund, M. Soljačić, C. Denz, D. A. Miller, and D. Psaltis, “Inference in artificial intelligence with deep optics and photonics,” Nature, vol. 588, no. 7836, pp. 39–47, 2020.
  * [5] X. Lin, Y. Rivenson, N. T. Yardimci, M. Veli, Y. Luo, M. Jarrahi, and A. Ozcan, “All-optical machine learning using diffractive deep neural networks,” Science, vol. 361, no. 6406, pp. 1004–1008, 2018.
  * [6] T. Zhou, X. Lin, J. Wu, Y. Chen, H. Xie, Y. Li, J. Fan, H. Wu, L. Fang, and Q. Dai, “Large-scale neuromorphic optoelectronic computing with a reconfigurable diffractive processing unit,” Nature Photonics, vol. 15, no. 5, pp. 367–373, 2021.
  * [7] Y. Zuo, B. Li, Y. Zhao, Y. Jiang, Y.-C. Chen, P. Chen, G.-B. Jo, J. Liu, and S. Du, “All-optical neural network with nonlinear activation functions,” Optica, vol. 6, no. 9, pp. 1132–1137, 2019.
  * [8] T. Wang, S.-Y. Ma, L. G. Wright, T. Onodera, B. Richard, and P. L. McMahon, “An optical neural network using less than 1 photon per multiplication,” arXiv preprint arXiv:2104.13467, 2021.
  * [9] Z. Gu, Y. Gao, and X. Liu, “Optronic convolutional neural networks of multi-layers with different functions executed in optics for image classification,” Optics Express, vol. 29, no. 4, pp. 5877–5889, 2021.
  * [10] M. Miscuglio, Z. Hu, S. Li, J. K. George, R. Capanna, H. Dalir, P. M. Bardet, P. Gupta, and V. J. Sorger, “Massively parallel amplitude-only fourier neural network,” Optica, vol. 7, no. 12, pp. 1812–1819, 2020.
  * [11] X. Porte, A. Skalli, N. Haghighi, S. Reitzenstein, J. A. Lott, and D. Brunner, “A complete, parallel and autonomous photonic neural network in a semiconductor multimode laser,” Journal of Physics: Photonics, vol. 3, no. 2, p. 024017, 2021.
  * [12] X. Xu, M. Tan, B. Corcoran, J. Wu, A. Boes, T. G. Nguyen, S. T. Chu, B. E. Little, D. G. Hicks, R. Morandotti, et al., “11 tops photonic convolutional accelerator for optical neural networks,” Nature, vol. 589, no. 7840, pp. 44–51, 2021.
  * [13] G. Mourgias-Alexandris, A. Tsakyridis, N. Passalis, A. Tefas, K. Vyrsokinos, and N. Pleros, “An all-optical neuron with sigmoid activation function,” Optics Express, vol. 27, no. 7, pp. 9620–9630, 2019.
  * [14] F. Duport, B. Schneider, A. Smerieri, M. Haelterman, and S. Massar, “All-optical reservoir computing,” Optics express, vol. 20, no. 20, pp. 22783–22795, 2012.
  * [15] F. Duport, A. Smerieri, A. Akrout, M. Haelterman, and S. Massar, “Fully analogue photonic reservoir computer,” Scientific Reports, vol. 6, no. 1, pp. 1–12, 2016.
  * [16] B. J. Shastri, M. A. Nahmias, A. N. Tait, A. W. Rodriguez, B. Wu, and P. R. Prucnal, “Spike processing with a graphene excitable laser,” Scientific Reports, vol. 6, no. 1, pp. 1–12, 2016.
  * [17] A. Dejonckheere, F. Duport, A. Smerieri, L. Fang, J.-L. Oudar, M. Haelterman, and S. Massar, “All-optical reservoir computer based on saturation of absorption,” Optics Express, vol. 22, no. 9, pp. 10868–10881, 2014.
  * [18] Y. Shen, N. C. Harris, S. Skirlo, M. Prabhu, T. Baehr-Jones, M. Hochberg, X. Sun, S. Zhao, H. Larochelle, D. Englund, et al., “Deep learning with coherent nanophotonic circuits,” Nature Photonics, vol. 11, no. 7, pp. 441–446, 2017.
  * [19] J. Feldmann, N. Youngblood, C. D. Wright, H. Bhaskaran, and W. H. Pernice, “All-optical spiking neurosynaptic networks with self-learning capabilities,” Nature, vol. 569, no. 7755, pp. 208–214, 2019.
  * [20] J. Feldmann, N. Youngblood, M. Karpov, H. Gehring, X. Li, M. Stappers, M. Le Gallo, X. Fu, A. Lukashchuk, A. S. Raja, et al., “Parallel convolutional processing using an integrated photonic tensor core,” Nature, vol. 589, no. 7840, pp. 52–58, 2021.
  * [21] F. Ashtiani, A. J. Geers, and F. Aflatouni, “Single-chip photonic deep neural network for instantaneous image classification,” arXiv preprint arXiv:2106.11747, 2021.
  * [22] S. Xu, J. Wang, H. Shu, Z. Zhang, S. Yi, B. Bai, X. Wang, J. Liu, and W. Zou, “Optical coherent dot-product chip for sophisticated deep learning regression,” arXiv preprint arXiv:2105.12122, 2021.
  * [23] B. Shi, N. Calabretta, and R. Stabile, “Inp photonic integrated multi-layer neural networks: Architecture and performance analysis,” APL Photonics, 2021.
  * [24] C. Wang, C. Langrock, A. Marandi, M. Jankowski, M. Zhang, B. Desiatov, M. M. Fejer, and M. Lončar, “Ultrahigh-efficiency wavelength conversion in nanophotonic periodically poled lithium niobate waveguides,” Optica, vol. 5, no. 11, pp. 1438–1441, 2018.
  * [25] M. Jankowski, C. Langrock, B. Desiatov, A. Marandi, C. Wang, M. Zhang, C. R. Phillips, M. Lončar, and M. Fejer, “Ultrabroadband nonlinear optics in nanophotonic periodically poled lithium niobate waveguides,” Optica, vol. 7, no. 1, pp. 40–46, 2020.
  * [26] Q. Guo, R. Sekine, L. Ledezma, R. Nehra, D. J. Dean, A. Roy, R. M. Gray, S. Jahani, and A. Marandi, “Femtojoule, femtosecond all-optical switching in lithium niobate nanophotonics,” arXiv preprint arXiv:2107.09906, 2021.
  * [27] L. Ledezma, R. Sekine, Q. Guo, R. Nehra, S. Jahani, and A. Marandi, “Intense optical parametric amplification in dispersion engineered nanophotonic lithium niobate waveguides,” arXiv preprint arXiv:2104.08262, 2021.
  * [28] D.-A. Clevert, T. Unterthiner, and S. Hochreiter, “Fast and accurate deep network learning by exponential linear units (elus),” arXiv preprint arXiv:1511.07289, 2015.
  * [29] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al., “Language models are few-shot learners,” arXiv preprint arXiv:2005.14165, 2020.
  * [30] M. Ono, M. Hata, M. Tsunekawa, K. Nozaki, H. Sumikura, H. Chiba, and M. Notomi, “Ultrafast and energy-efficient all-optical switching with graphene-loaded deep-subwavelength plasmonic waveguides,” Nature Photonics, vol. 14, no. 1, pp. 37–43, 2020.
  * [31] G. Grinblat, M. P. Nielsen, P. Dichtl, Y. Li, R. F. Oulton, and S. A. Maier, “Ultrafast sub–30-fs all-optical switching based on gallium phosphide,” Science advances, vol. 5, no. 6, p. eaaw3262, 2019.
  * [32] L. Deng, “The mnist database of handwritten digit images for machine learning research,” IEEE Signal Processing Magazine, vol. 29, no. 6, pp. 141–142, 2012.
  * [33] S. Bandyopadhyay, R. Hamerly, and D. Englund, “Hardware error correction for programmable photonics,” Optica, vol. 8, pp. 1247–1255, Oct 2021.
  * [34] X. Guo, T. D. Barrett, Z. M. Wang, and A. Lvovsky, “Backpropagation through nonlinear units for the all-optical training of neural networks,” Photonics Research, vol. 9, no. 3, pp. B71–B80, 2021.
  * [35] I. A. Williamson, T. W. Hughes, M. Minkov, B. Bartlett, S. Pai, and S. Fan, “Reprogrammable electro-optic nonlinear activation functions for optical neural networks,” IEEE Journal of Selected Topics in Quantum Electronics, vol. 26, no. 1, pp. 1–12, 2019.
  * [36] A. N. Tait, T. F. De Lima, E. Zhou, A. X. Wu, M. A. Nahmias, B. J. Shastri, and P. R. Prucnal, “Neuromorphic photonic networks using silicon photonic weight banks,” Scientific Reports, vol. 7, no. 1, pp. 1–10, 2017.
  * [37] J. Crnjanski, M. Krstić, A. Totović, N. Pleros, and D. Gvozdić, “Adaptive sigmoid-like and prelu activation functions for all-optical perceptron,” Optics Letters, vol. 46, no. 9, pp. 2003–2006, 2021.
  * [38] R. Amin, J. George, S. Sun, T. Ferreira de Lima, A. N. Tait, J. Khurgin, M. Miscuglio, B. J. Shastri, P. R. Prucnal, T. El-Ghazawi, et al., “Ito-based electro-absorption modulator for photonic neural activation function,” APL Materials, vol. 7, no. 8, p. 081112, 2019.
  * [39] C. Mesaritakis, A. Kapsalis, A. Bogris, and D. Syvridis, “Artificial neuron based on integrated semiconductor quantum dot mode-locked lasers,” Scientific Reports, vol. 6, no. 1, pp. 1–10, 2016.
  * [40] S. Oh, Y. Shi, J. Del Valle, P. Salev, Y. Lu, Z. Huang, Y. Kalcheim, I. K. Schuller, and D. Kuzum, “Energy-efficient mott activation neuron for full-hardware implementation of neural networks,” Nature Nanotechnology, vol. 16, no. 6, pp. 680–687, 2021.
  * [41] O. Krestinskaya, K. N. Salama, and A. P. James, “Learning in memristive neural network architectures using analog backpropagation circuits,” IEEE Transactions on Circuits and Systems I: Regular Papers, vol. 66, no. 2, pp. 719–732, 2018.
  * [42] Y. Huang, Z. Yang, J. Zhu, and T. T. Ye, “Analog circuit implementation of neurons with multiply-accumulate and relu functions,” in Proceedings of the 2020 on Great Lakes Symposium on VLSI, pp. 493–498, 2020.
  * [43] M. Giordano, G. Cristiano, K. Ishibashi, S. Ambrogio, H. Tsai, G. W. Burr, and P. Narayanan, “Analog-to-digital conversion with reconfigurable function mapping for neural networks activation function acceleration,” IEEE Journal on Emerging and Selected Topics in Circuits and Systems, vol. 9, no. 2, pp. 367–376, 2019.
  * [44] J. Choquette, W. Gandhi, O. Giroux, N. Stam, and R. Krashinsky, “Nvidia a100 tensor core gpu: Performance and innovation,” IEEE Micro, vol. 41, no. 2, pp. 29–35, 2021.
  * [45] Q. Xie, X. Lin, Y. Wang, S. Chen, M. J. Dousti, and M. Pedram, “Performance comparisons between 7-nm finfet and conventional bulk cmos standard cell libraries,” IEEE Transactions on Circuits and Systems II: Express Briefs, vol. 62, no. 8, pp. 761–765, 2015.
  * [46] M. Zhang, C. Wang, P. Kharel, D. Zhu, and M. Lončar, “Integrated lithium niobate electro-optic modulators: when performance meets scalability,” Optica, vol. 8, no. 5, pp. 652–667, 2021.
  * [47] J. Lu, A. Al Sayem, Z. Gong, J. B. Surya, C.-L. Zou, and H. X. Tang, “Ultralow-threshold thin-film lithium niobate optical parametric oscillator,” Optica, vol. 8, no. 4, pp. 539–544, 2021.
  * [48] Y. Yamamoto, K. Aihara, T. Leleu, K.-i. Kawarabayashi, S. Kako, M. Fejer, K. Inoue, and H. Takesue, “Coherent ising machines—optical neural networks operating at the quantum limit,” npj Quantum Information, vol. 3, no. 1, pp. 1–15, 2017.


11institutetext: University of Southern California, Los Angeles, CA, USA
11email<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
# An Intermediate Fusion ViT Enables Efficient Text-Image Alignment in
Diffusion Models

Zizhao Hu 11 Shaochong Jia 11 Mohammad Rostami 11

###### Abstract

Diffusion models have been widely used for conditional data cross-modal
generation tasks such as text-to-image and text-to-video. However, state-of-
the-art models still fail to align the generated visual concepts with high-
level semantics in a language such as object count, spatial relationship, etc.
We approach this problem from a multimodal data fusion perspective and
investigate how different fusion strategies can affect vision-language
alignment. We discover that compared to the widely used early fusion of
conditioning text in a pretrained image feature space, a specially designed
intermediate fusion can: (i) boost text-to-image alignment with improved
generation quality and (ii) improve training and inference efficiency by
reducing low-rank text-to-image attention calculations. We perform experiments
using a text-to-image generation task on the MS-COCO dataset. We compare our
intermediate fusion mechanism with the classic early fusion mechanism on two
common conditioning methods on a U-shaped ViT backbone. Our intermediate
fusion model achieves higher CLIP Score and lower FID, with $20\%$ reduced
FLOPs, and $50\%$ increased training speed compared against a strong U-ViT
baseline with an early fusion.

###### Keywords:

Diffusion models data fusion intermediate fusion

##  1 Introduction

Diffusion models [20] have emerged as a potent framework for generating high-
definition images [18, 6, 5, 17]. Most existing methods leverage prealigned
text embeddings such as CLIP (Dall-E [17], Dall-E 2 [16], Dall-E 3 [3], Stable
Diffusion [19], Stable Diffusion XL [14]) and use concatenation or cross-
attention to fuse the conditioning text embedding to the main image diffusion
features. The core issue of this approach is the inherit gap between image and
text representations in CLIP, plus the loss of specificity in certain
contexts. Therefore, CLIP struggles in abstract tasks such as counting, in
several types of fine-grained classification, and in out-of-distribution data
[15]. All of these discussed features are essential to be learned in the
diffusion process. Thus in the diffusion model, additional mechanisms are
needed to align the text embedding to the image diffusion task. The exact
location to fuse them should also reflect the natural information discrepancy
between these two data types. However, existing methods lack such
considerations. Some use simple concatenation or addition to fuse frozen text
embeddings at the diffusion input, while others use only a single cross-
attention layer to fuse the embeddings at all levels of the image diffusion
model. These approaches directly inherit the limitations of CLIP embeddings,
potentially damage the model’s performance on the diffusion learning task, and
introduce redundant text guidance with additional computing complexity.

To enable better and more efficient text-image alignment, we bring the
concepts from multimodal fusion to diffusion model conditioning and study how
different fusion techniques will impact the model’s performance on text-to-
image generation tasks. We investigate a specific type of ViT-based
segmentation backbones for the diffusion task. By examining the text-to-image
and image-to-image attention maps at different layers at all time stamps, we
discover that semantic information from captions provides more guidance at
bottleneck layers whereas fusing captions at earlier or later layers provides
minimum guidance. Based on this observation, we introduce an intermediate
fusion mechanism for ViT-based diffusion backbones. This method removes the
transformer blocks from both beginning and end layers and adds additional
text-only transformer layers. Further analytical experiments indicate that
this method improves the efficiency of text-image cross-attention mechanism.
Our method enhances the efficiency of the remaining cross-attentions near the
intermediate layers, leading to a better text-aligned, less interfered, and
efficient image diffusion model.

As a result, our proposed intermediate fusion enhanced backbone can generate
better quality and more text-aligned images, especially for high-level
semantics such as accurate object count, compound concepts, relationships
between multiple objects, out-of-distribution generations including rare
concepts and rare combinations of objects, etc.(see Figure 1 for a few
generated samples). Our model is also more efficient compared to the commonly
used early fusion mechanism in that it requires less computing, memory,
training time, and inference time.

Figure 1: Text-to-image data generation using our intermediate fusion
mechanism and the classic early fusion mechanism: images are generated based
on a text input using the corresponding mechanism in each row. The input text
contains a type of high-level semantics in each column. We observe that our
mechanism generates better quality samples that align better with the text
query while also being more efficient.

##  2 Background and Related Work

ViT-based flow-based model backbones have been explored recently, with
significant success in DiT [13] and U-ViT [1]. They bring several large-scale
applications in the text-guided generation domain (sora [9], stable diffusion
3 [19]) and multimodal generation domain (uni-diffuser [2]). These models
leverage different mechanisms to fuse the text with the generation modalities,
with simple concatenation and cross-attention being the two commonly used
mechanisms. Most existing works introduce the text as the pretrained embedding
features, which can be considered as early fusion in a diffusion model
setting. Recent work suggests removing text embedding from the last few layers
will reduce model FLOPs without affecting the model performance [13]. However,
this effect has not been studied systematically. As for text-image alignment,
most existing works try to improve the alignment from a training perspective
including finetuning with augmented data [11, 3], introducing additional
alignment guidance [23], etc. Quite differently, we approach addressing this
problem from an architectural perspective to enable better cross-modal fusion.
We briefly introduce the relevant background of our experiment setting.

###  2.1 Diffusion Models

Score-based diffusion models [21, 6] are a class of generative models that
learn to generate new samples from a target data distribution $p(x)$ by
approximating the score function using a neural network. The score function
provides a way to guide the generation of data without directly calculating
complex probability normalization. It also mitigates the risk of mode
collapse, encouraging a more diverse set of outputs that more accurately
reflects the complexity of the target distribution. The score function is
defined as the gradient of the log-probability density of the training data
points:

$S(x)=\nabla_{x}\log p(x).$ (1)

The training objective in score-based diffusion models is to minimize the
discrepancy between the true score function and its neural network
approximation. This is achieved by reducing the mean square error (MSE)
between the true score function, $S(\cdot)$, and the approximation provided by
the neural network, $s_{\theta}(\cdot)$:

$L(\theta)=\mathbb{E}[(S(x_{t})-s_{\theta}(x_{t}))^{2}].$ (2)

This modeling of the score function enables the model to generate the data
distribution from a noisy distribution by iteratively adjusting its parameters
in the direction of the gradient. This process involves taking small steps
along the gradient direction to optimize the score function approximation.

####  2.1.1 Training.

Given a dataset $\mathcal{D}$ of independently and identically distributed
(i.i.d.) real samples from the data distribution $p(x)$, the training process
of the score-based diffusion model involves two steps:

(i) Noise addition. Given a data sample $x$, we progressively inject Gaussian
noise over $T$ steps until a noisy observation $x_{T}$ of the input data is
generated. This process can be formalized as follows:

$x_{t}=\sqrt{1-\beta_{t}}x_{t-1}+\sqrt{\beta_{t}}\epsilon_{t},$ (3)

where $\epsilon_{t}\sim\mathcal{N}(0,I)$, $t=1,\ldots,T$, $x_{0}=x$, and
$\\{\beta_{t}\\}_{t=0}^{T}$ is a noise schedule.

(ii) Score function learning. The function $s_{\theta}(x_{t},t)$ approximates
the gradient of the log-density of the distribution $p(x_{t}|x)$. Let
$S_{t}=s_{\theta}(x_{t},t)$ be the approximated score function at time $t$,
and $G_{t}=\nabla_{x_{t}}\log p(x_{t}|x)$ be the gradient of the log
probability with respect to $x_{t}$. Then, the loss function
$\mathcal{L}(\theta)$ can be rewritten as:

$\mathcal{L}(\theta)=\mathbb{E}_{p_{data}(x)}\left[\frac{1}{2T}\sum_{t=0}^{T-1}\|S_{t}-G_{t}\|^{2}\right].$
(4)

####  2.1.2 Sampling.

The sampling (inference) process of the score-based diffusion model is a noise
addition process in a reverse direction:

(i) Initial noisy observation. We first initialize the noisy sample $x_{T}$
from a simple noise distribution, such as a standard Gaussian distribution.

(ii) Denoising process. We then apply the learned score function to denoise
$x_{T}$ over $t\in[T,1]$ steps to gradually remove noise and generate a
sample:

$x_{t-1}=\frac{1}{\sqrt{1-\beta_{t}}}(x_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\beta_{t}}}}s_{\theta}(x_{t},t))+\sqrt{\beta_{t}}\epsilon,$
(5)

###  2.2 Guided diffusion models

Guided generation of data involves creating new data samples, $x$, based on
given pair of data $(x_{0},y_{0})$. We can condition the score function on a
label $y$ to learn conditional generative models for class-specific data
generation. This objective can be done by leveraging the Bayes’ rule to
decompose the conditional score function [5]:

$\begin{split}S(x_{t}|y)&=\nabla_{x_{t}}\log p(x_{t}|y)=\nabla_{x_{t}}\log
p(y|x_{t})+\nabla_{x_{t}}\log p(x_{t})\end{split}$ (6)

Classifier guidance. By providing a weight term $\lambda$ to adjust the
balance between the unconditional score function and the classifier guidance
during the sampling phase, we can recast Eq. 6 as follows:

$\begin{split}S(x_{t}|y)=&\,\lambda\underbrace{\nabla_{x_{t}}\log
p(y|x_{t})}_{\text{classifier guidance}}+\underbrace{\nabla_{x_{t}}\log
p(x_{t})}_{\text{unconditional score function}}\end{split}$ (7)

where the first term can be learned by a classifier.

Classifier-free guidance (CFG). We can also model the unconditional score
function $\nabla_{x_{t}}\log p(x_{t})$ and the joint score function
$\nabla_{x_{t}}\log p(x_{t},y)$ simultaneously to substitute the classifier
guidance [7] for obtaining a trade-off between the quality and diversity of
samples. We replace $\lambda$ in Eq. 7 with $1+\omega$ to be consistent with
the range $[0,+\infty]$ of the guidance scale $\omega$ proposed in the
original paper:

$\begin{split}S(x_{t}|y)=&\,(1+\omega)\nabla_{x_{t}}\log
p(x_{t},y)-\omega\nabla_{x_{t}}\log p(x_{t})\end{split}$ (8)

The conditional generation in our study is achieved through classifier-free
guidance from caption $y$ to image $x$, while unconditional generation uses
the same approach with an empty caption embedding.

###  2.3 Latent Diffusion Models

Latent Diffusion Models (LDMs) [19] operate directly in the latent embedding
space of pretrained image features. By working in a lower-dimension latent
space, LDMs enable efficient generation and training. The formulation is
similar to vanilla diffusion models, except that a pretrained encoder
$\phi(\cdot)$ such as CLIP is used to map images to a latent space
$z=\phi{(x)}$ and the generated image is reconstructed from the denoised
latent variable $\hat{z}$ using the pretrained decoder
$\hat{x}=\theta{(\hat{z})}$. LDMs help enable efficient high-definition
sampling with less computational burden.

###  2.4 Conditioning Methods

Our major contribution is to enhance the current conditioning approaches by
replacing early fusion with intermediate fusion. In our empirical exploration,
we mainly focus on two common conditioning methods used by SOTA text-to-image
diffusion models:

(i) Concatenation. This approach concatenates the text embedding with the
image features and leverages the self-attention mechanism in subsequent layers
to fuse and guide image diffusion learning. See Figure 2 top row left.

(ii) Cross-attention. This approach avoids additional text tokens in the
intermediate layer by adding the text-to-image cross-attention to the image-
to-image self-attention at each fusion level. See Figure 2 top row right.

We only use these two common approaches to show that our method is applicable
and effective on different conditioning approaches. Although Other
conditioning methods have demonstrated effectiveness in recent text-to-image
diffusion models, testing all of them is beyond our resources and scope. We
contend that our approach can be seamlessly integrated to other conditioning,
provided that the dimensionality of internal representations remains uniform.

##  3 Proposed Methodology

Diffusion models are similar to hierarchical VAE decoders [22] in the aspect
of decoding intermediate latent space with a hierarchical generation process.
Such hierarchy leads to an observation that spatial and detailed information
at lower levels do not align with the higher-level semantics in image captions
that summarize the context. Our key intuition is that within the diffusion
process if text conditioning is forced to occur later at a layer primarily
concerned with higher-level semantics, more resources will be devoted to
achieving a match with high-level semantic content and better alignment will
be achieved. As a result, lower-level layers will suffer less from the
interference caused by the inherit gap between CLIP text and image
representations. We consider addressing the such interference using
intermediate fusion which centers around two core ideas: (i) additional
trainable text layers can pre-align text with the main diffusion model and
(ii) fusing these learnable embeddings at intermediate layers to simulate the
natural semantic information density of language and visual data.

###  3.1 Intermediate Fusion

In our study, we consider two competitive multimodal fusion strategies from a
text-to-image latent diffusion model perspective: (i) early fusion, where text
conditioning is fused with image at the beginning, and (ii) intermediate
fusion, where such fusion occurs at intermediate level. Most existing
diffusion models only use the early fusion strategy. In this method, both the
text and image embeddings are pretrained, frozen and concatenated at the input
level, as shown in Figure 2 bottom row left. To address the shortcomings of
this approach, we devise a diffusion backbone that has separate dedicated
trainable layers for each data. Then the layer outputs for text and image will
be concatenated at intermediate level, as shown in Figure 2 bottom row right.

###  3.2 Diffusion Backbone Model

Evidence suggests that under a diffusion model setting, segmentation networks
with long skip connections are essential to the efficient learning of discrete
time ODE [8]. When long skip connections are used, distant network blocks can
be connected to aggregate long-distant information and alleviate vanishing
gradient. For this reason, we choose U-ViT-Small [1] as our strong baseline.
The baseline model uses self-attention with early fusion (Figure 2 entire left
column). On top of this, we study three other settings, namely cross-attention
with intermediate fusion (Figure 2 entire right column), self-attention with
intermediate fusion, cross-attention with early fusion (Figure 2 diagonal
combinations). These settings are constructed by only changing the backbone
models without modifying the training and inference process.

Figure 2: The two conditioning methods (top row) and two fusion methods
(bottom row) in our experiments. The baseline is a 13-layer ViT model with the
same skip-connection mechanisms introduced in U-ViT. For fusion comparison, we
carefully designed $N_{\text{joint}}$ to ensure same depth and number of
parameters for the image branch. In the experiments we use
$N_{\text{image}}=4$ and $N_{\text{text}}=1$. For conditioning comparison, we
switch the attention type and skip connections with other blocks intact.

###  3.3 Latent Diffusion

We use a stable diffusion KL-based autoencoder to encode an input image into
the latent space and decode the denoised latent space representation back to
the input image space. For text embeddings, we use the CLIP embedding with
ViT-L-14. These models are frozen during diffusion model training.

##  4 Experiments

###  4.1 Dataset and Training Settings

In our experiments, we use the MSCOCO [10] train and validation datasets to
train and evaluate the performance of our model. For our training
configuration, we train all models for 1 million steps and use a batch size of
256. We use the AdamW optimizer, with a learning rate of 0.0002, weight decay
of 0.03, and beta parameters set to (0.9, 0.9). We incorporate a warm-up phase
of 5000 steps to adjust the learning rate. The ViT model takes image features
with a channel of 4, both spatial dimensions of 32, and an image patch size of
2. All attention mechanisms use an embedding dimension of 512, 8 attention
heads. CLIP embedding has 77 tokens each with a dimension of 768, and is
transformed to 512 dimension using a linear layer to align with the
transformer input. For classifier-free guidance, we use a probability of 0.1
for unconditional training. For intermediate fusion, we use
$N_{\text{image}}=4$ and $N_{\text{text}}=1$.

###  4.2 Evaluation Metrics

We deploy two evaluation methods, quantitative evaluation and human
evaluation. Recent researches discover that FID cannot perfectly reflect the
human preference for generation quality [4, 12], and CLIP Score cannot focus
on specific alignment aspects such as matching object count [15]. Therefore,
we also incorporate human evaluation for generation from specially designed
prompts.

(i) Quantitative evaluation. We use FID, and CLIP Score as our quantitative
metrics. To generate the score we select 30000 captions from the MSCOCO
validation set and the corresponding generated images from our text-to-image
models. For CLIP Score we use the CLIP version CLIP-ViT-L-14.

(ii) Human evaluation - object count. We choose a challenging generation
aspect even for most of the foundation text-to-image models - matching object
count, where we require the model to generate the same amount of objects as
described in the prompts. We select four objects - bus(es), sheep,
person(people), and apple(s). These four are selected since they represent 4
different plural forms and 4 categories(human, animal, fruit, human-made
object). We use 5 words of count - a(an), two, three, four, five. Since larger
numbers are rare in MS-COCO training captions, we restrict our study to small
numbers. We generate 10 images for each object-count pair(20 pairs) and let
evaluators count the number of target objects in the generated image. Then we
use the average error and average match ratio to evaluate the performance:

$\text{Average\
Error}=\frac{\sum_{i=1}^{n}|\text{count}_{i,\text{human}}-\text{count}_{i,\text{prompt}}|}{n}$
(9) $\text{Average\ Match\
Ratio}=\frac{\sum_{i=1}^{n}\mathbb{I}(\text{count}_{i,\text{human}}=\text{count}_{i,\text{prompt}})}{n}$
(10)

(iii)Human evaluation - preference score.  In addition, we ask 5 evaluators to
provide a preference ranking for 1 to 4 on the overall quality of images
generated by each model from random captions in the evaluation set. We use the
same random seed and prompts for all models and provide the prompts and
generated images to the evaluators. The human evaluators are allowed to skip
any group of samples if the ranks are difficult to call. 100 captions are
evaluated by 5 evaluators, with a maximum of 500 scores for each setting. We
assign 4, 3, 2, and 1 scores to rank 1, 2, 3, and 4 respectively, and
calculate the average score for each model setting.

###  4.3 Results

We selected U-ViT [1] as our baseline model, since it has the best
reproducible FID score among dedicated diffusion models trained on the MS-COCO
dataset to our knowledge. We first compare the performance between these 4
experiment settings in Table 1.

Table 1: Comparative results on text-to-image Generation and alignment metrics. The baseline U-ViT-Small correspond to the first settings. FID is evaluated on generated images from 30K captions from the evaluation set. CLIP Score is calculated on 30K generated samples and ground truth caption pairs using CLIP-ViT-L-14. For training speed, an iteration(iters) is a full forward-backward pass on an RTX-4090 GPU with a mini-batch size of 256. GFLOPs are calculated on a single forward pass of the model at one timestamp. Conditioning Method |  Fusion Type |  FID-30K $\downarrow$ |  CLIP Score $\uparrow$ |  Params |  Training iters/s |  GFLOPs
---|---|---|---|---|---|---
Concat |  Early(baseline) |  5.98 |  0.584 |  45M |  1.81 |  29.56
|  Intermediate |  5.77 |  0.588 |  48M |  2.31 |  25.84
Cross-attn |  Early |  6.48 |  0.575 |  45M |  2.54 |  23.82
|  Intermediate |  5.68 |  0.588 |  48M |  2.74 |  23.66

In both conditioning methods, intermediate fusion shows lower FID and higher
CLIP Scores than its early fusion counterpart. Compared to baseline setup, all
other setups train faster and requires smaller GPU usage, even though some
have more parameters. This is due to the reduced attention calculation. Next,
we visualize the FID (Figure 3, left), CLIP Score (Figure 3, middle) during
training, and FID vs. CLIP Score at different CFG scales (Figure 3, right).

From the results so far, We find that intermediate fusion settings improve all
quantitative performance compared to their early fusion counterparts by a
large margin. Among all four settings, an intermediate fusion with cross-
attention has the best FID, CLIP score, lowest GFLOPs, and fastest training.

Figure 3: Evaluation during training and FID-30K vs CLIP Score at different
CFG scales. Intermediate fusion settings show improved generation quality and
text-image alignment compared to their early fusion counterparts. CLIP Score
is measured on 30K pairs using CLIP-ViT-L-14.

We select 12 random captions and generate images with a CFG scale of 3, which
is an elbow point in the FID vs. CLIP Score curve. We show the baseline (top)
compared with our best setting (bottom) in Figure 4.

Figure 4: Comparison between the best intermediate fusion model(bottom) and
the baseline model(top) across 12 different prompts. (Best viewed when zoomed-
in.)

We then show our best-performing setup against several foundation models and
dedicated models (Table 2). Our model can reach the best text-image alignment
performance and comparable image quality to all models with a relatively small
model size.

Table 2: Performance of transformer-based text-to-image diffusion models. Model |  Fusion Type |  Fusion Path |  FID-30K $\downarrow$ |  CLIP Score $\uparrow$ |  GFLOPs
---|---|---|---|---|---
Foundation Models Zero-shot on MS-COCO
DALLE 3 |  Early |  Cross-attn |  - |  0.320 |  -
Imagen |  Early |  Cross-attn |  7.27 |  $\sim$0.29 |  -
Stable Diffusion |  Early |  Cross-attn |  8.59 |  0.325 |  -
SDXL |  Early |  Cross-attn |  |  0.305 |  -
Dedicated Models Trained on MS-COCO
U-ViT-Small |  Early |  Concat |  5.98 |  0.584 |  29.56
IF-U-ViT(ours) |  Intermediate |  Cross-attn |  5.68 |  0.588 |  23.66

###  4.4 Human Evaluation

Object count. The results are shown in Figure 5. In the left four figures, for
18 out of 20 object-count pairs, the intermediate concatenation generates
objects with more human-aligned count compared to the early concatenation. For
14 out of 20 object-count pairs, the intermediate cross-attention generates
objects with more or equal human-aligned count compared to the early cross-
attention. In the top right figure, the average error of intermediate fusion
is consistently lower than the early fusion counterparts. In the bottom right
figure, the average match ratio of intermediate fusion is consistently higher
or on par with the early fusion counterparts. All of the above results show
that intermediate fusion improves the count alignment in the generation
regardless of the conditioning method.

Figure 5: Human evaluation on object count. Lighter colors represent early
fusion, while darker colors represent intermediate fusion. The left four
figures are the average error given different ground truth counts, where
x-axis is the ground truth. Each figure corresponds to an object. The right
top figure is the average error across all counts for different objects. The
bottom right figure is the average percentage of exact matches for each
object. The plots indicated lower average count errors and higher matching
counts of intermediate fusion.

Preference score The results are shown in Figure 6. 287 scores out of 500
expected scores are collected after removing invalid scores and those are too
difficult to call by the human evaluators. The score is consistent with our
FID and CLIP Score evaluation, with intermediate cross-attention achieving the
highest score, and intermediate concatenation coming second. All intermediate
fusion settings outperform their early fusion counterparts.

Figure 6: Human evaluation on general quality of generation.

###  4.5 Ablations

Our intermediate fusion includes two atomic methods: 1. additional trainable
transformer block for text embeddings during the image diffusion training 2.
Only fuse these embeddings at the middle layers of a diffusion backbone. We
study each of the method’s contributions to the FID, and CLIP Score
separately. In Figure 7, we show that fusing the text embedding only in the
middle of the diffusion backbone is the major source of FID improvement. This
is expected since this setting has image-only skip connections that can
maintain spatial consistency information at the upsampling layers. But this
will negatively impact the CLIP Score. Adding text transformers learns more
aligned text embeddings, which is the major source of improved CLIP Score, but
this will impact the model efficiency in terms of increased FLOPs and reduced
training speed. When the two methods are combined, we achieve improvements in
all metrics. We see that these two atomic methods compensate for each other’s
weaknesses while maintaining their respective advantages.

Figure 7: Ablation Study. We show the effect of adding trainable text
transformer layers and fusing the embedding in the middle separately, together
with the joint effect(Both) and baseline(None, early cross-attention).

###  4.6 Analysis

Layer-wise Attention Maps. To better visualize the text-image alignment across
the model layers, we analyze the average attention maps of all timesteps
during diffusion process.In Figure 8 we show the comparison of early and
intermediate fusion. Attention maps of early fusion models provide a valuable
observation.The text-to-image attention maps in both early and late layers
indicate a more uniformly distributed pattern than intermediate layers,
suggesting that text guidance is less focused and effective in early and late
layers. Besides this observation, the early and late layers attend more to the
border of the latent image due to the padded convolution encoding of images in
the autoencoder model. To reduce the influence of such padded borders, we
removed the border so that the later analysis can reflect the true image
semantic guidance.

Rank Analysis on Adjusted Attention Map  To quantify the influence of text
guidance on attention map image features, we conducted SVD on the attention
matrix and analyzed their rank property. We see that early fusion models have
relatively low-rank attention maps with smaller singular values at all layers,
especially the layers away from the middle. The intermediate fusion on the
other hand has high singular values. The analysis indicates that the early
fusion introduces low-efficiency text-to-image attention in early and end
layers of a ViT diffusion backbone, whereas most of the text information is
fused around the bottleneck. This justifies our presumption of inconsistent
alignment at lower levels of early fusion. Additionally, the comparison of
quantified result proves that the elimination of the early and end fused
layers never hurts the effectiveness of guidance. It instead boosted the
guidance. Thus, by forcing the fusion to occur at a later stage, we can
potentially improve model efficiency without damaging the semantic control of
text. This observation aligns with the experiment results.

Figure 8: Attention maps and singular value analysis. For each setting, the
first row is the image-to-image attention, the second row is the text-to-image
attention, and the third row is the singular values of the first 10 orders
from the text-to-image attention maps. An intermediate fusion removes the low-
information text-to-image attention calculations at the early and late levels.
This reduces the interference between image and text at these levels and
improves the information capacity of the text-to-image attention at the middle
layers. The combined effect results in better high-level semantic alignment
and low-level spatial features’ integrity.

##  5 Limitations

####  5.0.1 Other Conditioning Methods.

The main focus of this paper is to investigate the influence of intermediate
fusion provided by a new architecture backbone, in terms of text-image
alignment, generation quality, and computational efficiency compared to early
fusion models. Admittedly, early fusion covers a large number of conditioning
strategies other than concatenation and cross-attention, and some methods such
as the Adaptive Layer Norm used by DiT are not explored in this paper. We
reasonably argue that the intermediate fusion can be transferred with ease to
other unexplored conditioning strategies, since the approach resolved the
issue of less efficient attention caused by joining image and text at input
level.

Model Parameters and Hyperparameters. Our model contains two hyperparameters
$N_{\text{image}}$ and $N_{\text{text}}$(in latter discussion, a setting is
denoted as [$N_{\text{image}}$, $N_{\text{text}}$]). In our experiments, we
test one intermediate fusion setting([1,4]) in our main experiments and three
settings([1,4],[0,4], and [1,0]) in ablation studies. However, how different
architectural design will impact the results remains unexplored. This likely
depends on the pretrained autoencoder, pretrained text embedding, embedding
size, dataset, conditioning method, and other factors. Moreover, we logically
contend that our approach can be applied to scaled-up foundation models, and
will likely improve the performance with considerate design choices.

##  6 Conclusion

In this study, we presented an effective architecture for enhancing text-to-
image diffusion models by leveraging an intermediate fusion mechanism for
conditioning strategies. Our extensive experiments and analyses on the MS-COCO
dataset demonstrate that this method significantly outperforms traditional
early fusion techniques in aligning visual concepts with the high-level
semantics of language. By integrating trainable text embeddings in the middle
layers of a U-ViT backbone, our approach not only boosts the quality and text
alignment of generated images but also enhances the efficiency of the training
and inference. This improvement is quantitatively supported by improved CLIP
Scores and reduced FID values, alongside qualitative assessments through human
evaluations.

More generally, our findings suggest that the placement and integration of
text embeddings within diffusion models play a critical role in the overall
performance and efficiency of text-to-image generation tasks. By adopting an
intermediate fusion strategy, we successfully alleviate the prevalent issue of
misalignment between textual semantics and generated visual content. This
provides a direction for large foundation models to a more efficient and text-
aligned design.

## References

  * [1] Bao, F., Nie, S., Xue, K., Cao, Y., Li, C., Su, H., Zhu, J.: All are worth words: A vit backbone for diffusion models (2023)
  * [2] Bao, F., Nie, S., Xue, K., Li, C., Pu, S., Wang, Y., Yue, G., Cao, Y., Su, H., Zhu, J.: One transformer fits all distributions in multi-modal diffusion at scale. arXiv preprint arXiv:2303.06555 (2023)
  * [3] Betker, J., Goh, G., Jing, L., TimBrooks, ., Wang, J., Li, L., LongOuyang, ., JuntangZhuang, ., JoyceLee, ., YufeiGuo, ., WesamManassra, ., PrafullaDhariwal, ., CaseyChu, ., YunxinJiao, ., Ramesh, A.: Improving image generation with better captions. https://api.semanticscholar.org/CorpusID:264403242
  * [4] Borji, A.: Pros and cons of gan evaluation measures. Computer vision and image understanding 179, 41–65 (2019)
  * [5] Dhariwal, P., Nichol, A.: Diffusion models beat gans on image synthesis (2021)
  * [6] Ho, J., Jain, A., Abbeel, P.: Denoising diffusion probabilistic models (2020)
  * [7] Ho, J., Salimans, T.: Classifier-free diffusion guidance (2022)
  * [8] Huang, Z., Zhou, P., Yan, S., Lin, L.: Scalelong: Towards more stable training of diffusion model via scaling network long skip connection. Advances in Neural Information Processing Systems 36 (2024)
  * [9] Karaarslan, E., Aydın, Ö.: Generate impressive videos with text instructions: A review of openai sora, stable diffusion, lumiere and comparable models (2024)
  * [10] Lin, T.Y., Maire, M., Belongie, S., Bourdev, L., Girshick, R., Hays, J., Perona, P., Ramanan, D., Zitnick, C.L., Dollár, P.: Microsoft coco: Common objects in context (2015)
  * [11] Paiss, R., Ephrat, A., Tov, O., Zada, S., Mosseri, I., Irani, M., Dekel, T.: Teaching clip to count to ten (2023)
  * [12] Parmar, G., Zhang, R., Zhu, J.Y.: On aliased resizing and surprising subtleties in gan evaluation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 11410–11420 (2022)
  * [13] Peebles, W., Xie, S.: Scalable diffusion models with transformers. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 4195–4205 (2023)
  * [14] Podell, D., English, Z., Lacey, K., Blattmann, A., Dockhorn, T., Müller, J., Penna, J., Rombach, R.: Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952 (2023)
  * [15] Radford, A., Kim, J.W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al.: Learning transferable visual models from natural language supervision. In: International conference on machine learning. pp. 8748–8763. PMLR (2021)
  * [16] Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., Chen, M.: Hierarchical text-conditional image generation with clip latents (2022)
  * [17] Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Radford, A., Chen, M., Sutskever, I.: Zero-shot text-to-image generation. In: International Conference on Machine Learning. pp. 8821–8831. PMLR (2021)
  * [18] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., Ommer, B.: High-resolution image synthesis with latent diffusion models. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 10684–10695 (2022)
  * [19] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., Ommer, B.: High-resolution image synthesis with latent diffusion models (2022)
  * [20] Sohl-Dickstein, J., Weiss, E.A., Maheswaranathan, N., Ganguli, S.: Deep unsupervised learning using nonequilibrium thermodynamics (2015)
  * [21] Song, Y., Sohl-Dickstein, J., Kingma, D.P., Kumar, A., Ermon, S., Poole, B.: Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 (2020)
  * [22] Vahdat, A., Kautz, J.: Nvae: A deep hierarchical variational autoencoder. Advances in neural information processing systems 33, 19667–19679 (2020)
  * [23] Wu, X., Sun, K., Zhu, F., Zhao, R., Li, H.: Human preference score: Better aligning text-to-image models with human preference (2023)

##  Appendix 0.A Appendix

This appendix provides additional details of the experiments and additional
results discussed in the main paper.

###  0.A.1 Detailed Experimental Setup

Provide detailed descriptions of the experimental setup that were not included
in the main manuscript due to space constraints. Include any specific
settings, parameters, or configurations used in your experiments.

Table 3: Detailed Configuration of the models.

Configuration Aspect | Parameter | Value
---|---|---
General Settings | Random Seed | 1234
| Tensor Dimensions | (4, 32, 32)
Model Architecture | Autoencoder Scaling Factor | 0.23010
Training | Number of Steps | 1,000,000
| Batch Size | 256
Optimization | Optimizer Type | AdamW
| Learning Rate | 0.0002
| Weight Decay | 0.03
| Momentum (Betas) | (0.9, 0.9)
Learning Schedule | Scheduler Type | Linear
| Warm-up Steps | 5,000
Network | Network Type | ViT with skip-connections
| Image Feature Size | 32
| Input Channels | 4
| Patch Size | 2
| Embedding Dimension | 512
| Network Depth | 12
| Number of Attention Heads | 8
| MLP Ratio | 4
Data Configuration | Dataset | MSCOCO-256x256
| CFG Unconditional Probability | 0.1
Sampling Strategy | Steps per Sample | 50
| Total Samples | 30,000
| Mini-batch Size | 50
| Sampling CFG Scale | 3.0
| Eval CFG Scale | 1.0

###  0.A.2 Additional Results

We present more text-to-image generations from 4 different models here. We use
the same prompt randomly selected from the MSCOCO evaluation set with a
classifier-free guidance level 3.

Figure 9: Early fusion + Concatenation(U-ViT-Small) Figure 10: Intermediate
fusion + concatenation Figure 11: Early fusion + cross-attention Figure 12:
Intermediate fusion + cross-attention

###  0.A.3 Human Evaluation Setup

We sent the following forms to 5 evaluators, together with 30 groups of
generated samples for object count, and 100 groups of generated samples for
preference ranking. The evaluators are not aware of the purpose of this
project and the model identification throughout the evaluation. The model
order is not randomized, so it’s possible there are biases passed to later
groups based on seen groups.

Table 4: Human evaluation for object count form. ID |  Question | Number of Objects Identified
---|---|---
1 |  How many {Object} can you see in the image? | 3
Table 5: Human evaluation for preference ranking form. The question is “Rank the quality of the generation given {Prompt}.” ID | Image 1 | Image 2 | Image 3 | Image 4 | Skipped (Y/N)
---|---|---|---|---|---
1 | 1 | 2 | 3 | 4 | N

###  0.A.4 Code Availability

The code will be released on GitHub.