Machine-learning-accelerated Bose-Einstein condensation

Recently, researchers have begun applying machine learning techniques to atomic physics experiments, e.g., to enhance data processing for imaging [1, 2, 3, 4, 5], determine the ground state and dynamics of many-body systems [6, 7], or to identify phases and phase transitions [8, 9, 10]. One promising practical application of machine learning to atomic physics is in the optimization of control sequences with many parameters and nonlinear dynamics [11, 12, 13, 14, 15, 16], and in particular to one of the workhorses of atomic physics, Bose-Einstein condensates (BECs) [11, 13, 14, 15, 16].

With few exceptions [17], experiments on BECs end with a destructive measurement, which requires repeated BEC preparation. Approaches to increase the BEC production rate, and associated signal-to-noise ratio of the experiments, have generally relied heavily on hardware improvements [18, 19, 20, 21, 22], or used atomic species with narrower optical transitions [18, 21, 22] than offered by the most widely utilized alkali atoms. For alkali atoms, the tight confinement of atom-chip magnetic traps has enabled fast evaporation sequences, with a complex multi-layer atom-chip achieving BEC preparation times of $850\ \mathrm{ms}$ for $4\times 10^{4}$ atoms [19]. Non-alkali atoms featuring narrow optical transitions can be used to reach lower temperatures in narrow-line MOTs [18, 21, 22]. That approach, combined with a dynamically tunable optical dipole trap, has recently been used to prepare BECs of $2\times 10^{4}$ erbium atoms in under $700\ \mathrm{ms}$ [22].

In this Letter, we demonstrate a complementary approach where, in a simple experimental setup with a broad-line MOT for a standard alkali atom, machine learning is leveraged to optimize a complex nonlinear laser and evaporative cooling process to quantum degeneracy. Controlling a sequence with up to 55 interdependent experimental parameters, Bayesian optimization [23, 11, 12] finds parameter values which cool a gas from room temperature into the quantum degenerate regime in $575\ \mathrm{ms}$, creating a BEC containing $N_{\mathrm{BEC}}=2.8\times 10^{3}$ atoms. To our knowledge, this is the fastest BEC creation to date. We identify some of the physical strategies discovered by the algorithm, and also investigate how the choice of cost function impacts the trade-off between final atom number and the purity of the created BEC.

Our apparatus employs only a single MOT directly loaded from a ${}^{87}\mathrm{Rb}$ background vapor, a crossed optical dipole trap, and two Raman cooling beams as depicted in Fig. 1(a). No Zeeman slower, two-dimensional MOT, atom chip [19], dynamic trap shaping [21], or strobing [24, 22] are necessary. Using Raman cooling in a crossed optical dipole trap (cODT), a method that can reach very high phase-space density and even condensation [25], the algorithm achieves a cooling slope of 16 orders of magnitude improvement in phase space density (PSD) per order of magnitude in atom loss ($\gamma=16$) up to the threshold to quantum degeneracy. This is significantly better than the $\gamma=7$ value we could obtain with extensive manual optimization under similar conditions [25].

Atomic physics methods.—The Raman cooling implementation used in this work is similar to that of Ref. [25]. Cooling proceeds in a cODT formed by intersecting two non-interfering $1064\ \mathrm{nm}$ beams, one horizontal and one vertical, see Fig. 1(a). Two $795$-nm beams drive the Raman cooling: the optical pumping beam and the Raman coupling beam. Raman cooling [26] provides sub-Doppler cooling by driving velocity-selective Raman transitions between hyperfine states, here the $|F=2,m_{F}=-2\rangle$ and $|2,-1\rangle$ states of ${}^{87}\mathrm{Rb}$ [25]. The Raman transitions are non-dissipative so entropy is removed from the atomic gas in the form of spontaneously scattered photons as atoms are optically pumped back to the dark state $|2,-2\rangle$. Light-assisted collisions, which typically prohibit laser cooling at high atomic densities, are suppressed by detuning the optical pumping light $4.33\ \mathrm{GHz}$ to the red of the $D_{1}$ $F=2\rightarrow F^{\prime}=2^{\prime}$ transition, where a local minimum of light-induced loss was observed [25].

The cooling dynamics are controlled via five actuators: (i) the horizontal $P_{y}$ and (ii) vertical $P_{z}$ trap beam powers which set the trap depth and frequencies, (iii) the Raman coupling beam power $P_{\mathrm{R}}$ which tunes the Raman rate, (iv) the power $P_{\mathrm{p}}$ of the optical pumping beam which sets the optical-pumping rate (and also Raman rate), and (v) the magnetic field $B_{z}$ which adjusts the resonant velocity class for the Raman transition. The cooling procedure is divided into stages during which the controls are linearly ramped, with the endpoints of each ramp constituting the optimization parameters.

Optimization scheme.—The optimization problem can be formulated as the minimization of a cost function $C$, which maps a set of parameter values $\mathbf{X}\in\mathbb{R}^{M}$ to a corresponding cost value $C(\mathbf{X})\in\mathbb{R}$, where $M$ is the number of optimization parameters. The cost $C$ quantifies the results, and is generally a priori unknown, but can be extracted from measurements. Bayesian optimization is well-suited for this type of problem as it can tolerate noise in the measured cost and typically requires testing fewer values of $\mathbf{X}$ than other optimization methods [11, 12, 13, 14, 15, 16].

Bayesian optimization begins with collecting a training dataset by experimentally measuring the cost $C_{\mathrm{m}}(\mathbf{X}_{i})$ for various values of sets of parameter values $\mathbf{X}_{i}$. The $\mathbf{X}_{i}$ used to construct the training dataset are chosen by a training algorithm, which can implement another optimization algorithm or can select $\mathbf{X}_{i}$ randomly. A model of the cost function is then fit to the training dataset which approximates the unknown true cost function $C(\mathbf{X})$. Although Bayesian optimization typically uses a Gaussian process for its model [23], the present work uses neural networks [12, 27], which were chosen for their significantly faster fitting time for our typical number of optimization parameters. Once the model is fit, a standard numerical optimization algorithm is applied to the modeled cost function $C_{\mathrm{p}}(\mathbf{X})$ to determine which value $\mathbf{X}_{i+1}$ for the next iteration is predicted to yield the minimal cost, as depicted in Fig. 1(c). Optionally this numerical optimization can be constrained to a trust region (a smaller volume of parameter space centered around the $\mathbf{X}_{i}$ which yielded the best cost measured thus far). The predicted optimal value $\mathbf{X}_{i+1}$ is then tested by experimentally measuring the corresponding cost $C_{\mathrm{m}}(\mathbf{X}_{i+1})$. The next iteration begins by retraining the model with the new result, and making a new prediction for the optimal value of $\mathbf{X}$ with the updated model. The algorithm iterates until it reaches a termination criterion, such as a set maximum number of iterations, or a set number of consecutive iterations that fail to return better results. All optimization in this work was performed with the open-source packages M-LOOP [11, 12] to implement the Bayesian optimization and Labscript [28] for experimental control. Additional implementation details are included in Appendix A.

Cost function.—Since the optimization transitions the gas from the classical into the quantum degenerate regime, the final state of the gas depends strongly on how the cost function is chosen as a combination of the two experimentally accessible parameters: atom number $N$ and temperature $T$. The classical phase space density $\mathrm{PSD}_{\mathrm{c}}$ is defined as $\mathrm{PSD}_{\mathrm{c}}\equiv n_{\mathrm{cp}}\lambda_{\mathrm{dB}}^{3}$, where $\lambda_{\mathrm{dB}}$ is the thermal de Broglie wavelength and $n_{\mathrm{cp}}$ is the calculated peak number density neglecting bosonic statistics (See Appendix B for calculation details). The value of $\mathrm{PSD}_{\mathrm{c}}$ is nearly equal to the true $\mathrm{PSD}$ when $\mathrm{PSD}\ll 1$, while at the threshold to condensation, $\mathrm{PSD}_{\mathrm{c}}\sim 1$. Since the temperature $T$ is more difficult to determine in the quantum degenerate regime, and also requires a fit to the data with potential convergence problems, we instead measure $N$ and the peak optical depth $\mathrm{OD}$ in an absorption image. Generally ensembles with larger $\mathrm{PSD}_{\mathrm{c}}$ have a larger atom number $N$ and less expansion energy, which leads to a larger peak optical depth $\mathrm{OD}$ for a given $N$. Guided by this, we explored cost functions of the form

where $f(N/N_{1})$ is a smoothed Heaviside step function with $N_{1}$ chosen near the detection noise floor (see Appendix A). The parameter $\alpha$ in the cost function tunes the trade-off between optimizing for larger atom number or lower temperature. For a pure BEC after sufficient time-of-flight (TOF) expansion, $\lvert C/f\rvert$ scales as $(N_{\mathrm{BEC}})^{\alpha}$ (see Appendix C). For a thermal cloud, $\lvert C/f\rvert$ is proportional to $\mathrm{PSD}_{\mathrm{c}}$ when $\alpha=-1/5$, although that value of $\alpha$ is unsuitable for condensation as increasing the atom number in the BEC requires $\alpha>0$.

Optimization procedure.—The sequence begins with a separately optimized $99$-ms long MOT loading and compression period. The trap beam powers are ramped to their initial Raman cooling values during the last $10\ \mathrm{ms}$ of the MOT compression and then the magnetic field is adjusted to its initial Raman cooling value in $1\ \mathrm{ms}$, at which point the horizontal dipole trap holds typically $N=2.7\times 10^{5}$ atoms. We then added $100$-ms stages of Raman cooling one by one and optimized them individually. After five stages, the algorithm tended to turn off the Raman cooling by turning down $P_{\mathrm{p}}$ or $P_{\mathrm{R}}$ or by tuning the magnetic field $B_{z}$ such that the Raman transition became off-resonant. We then added up to six shorter $30$-ms long stages in which the optical pumping and Raman coupling beams were turned off, and the algorithm performed evaporative cooling. Due to the reduced number of parameters, we were able to optimize the evaporation stages simultaneously, which produced a BEC. Subsequently we shortened the Raman cooling and evaporation stages with parameter values fixed until only a small and impure BEC was produced, and then ran a global reoptimization. In this global optimization stage, all 42 of the Raman cooling and evaporation parameters were reoptimized simultaneously using the previous values as the initial guess for $\mathbf{X}$. Often a trust region set to one tenth of the allowed range for each parameter was used. This kept the optimizer focused in regions of parameter space which produced a measurable signal, as adjusting even a single parameter too far would often result in the loss of all atoms. We repeated this sequence shortening and reoptimization procedure until the algorithm failed to find parameters that could produce sufficiently pure BECs.

The required beam powers generally varied over several orders of magnitude, so the logarithm of their powers were used as entries in $\mathbf{X}$, while the magnetic-field control parameter $B_{z}$ was kept a linear parameter. A feedforward adjustment was included in the $B_{z}$ control values to account for the light shift of the $|2,-1\rangle$ state by the optical pumping beam. We averaged over five repetitions of the experiment for each set of parameter values tested. The number of iterations per optimization varied but was typically ${\sim}1000$ (including the initial training), and required several hours, both for the single-stage optimizations and the full-sequence optimizations. A simpler optimization procedure was also attempted which did not involve optimizations of individual stages. Instead the sequence was divided into ten $100\ \mathrm{ms}$ stages and all 55 parameters were optimized from scratch simultaneously. That approach combined with the shortening and reoptimizing procedure successfully produced a similar BEC, albeit in slightly longer time ($650\ \mathrm{ms}$ vs $575\ \mathrm{ms}$), possibly due to the optimization becoming trapped in a local optimum (see Appendix A for further discussion).

Results and physical interpretation.—The best discovered $575$-ms long control sequence and corresponding results are depicted in Fig. 2 and Fig. 3. Notably, the algorithm discovered gray molasses [29, 30] in the MOT phase, which it applies at the end of the compression sequence. This outperforms the bright molasses [31, 32] that was previously used in the manually optimized compression sequence, with the gray molasses loading a similar number of atoms ten times faster. After the MOT loading stage and transfer into the cODT, five ${\sim}63$-ms long stages of Raman cooling follow, and then the optical pumping and Raman beams are ramped off, followed by six ${\sim}27$-ms long evaporation stages. As observed in previous work [12, 13, 14], the ramps produced by Bayesian optimization are non-monotonic and appear non-intuitive, but they outperform the routines we found by manual optimization. A reason for the non-monotonic waveforms may be that the cost function includes many local minima. The optimization can settle into any one of these local optima randomly, and produce complex but specific waveforms, as observed in Ref. [12]. Despite the non-monotonic ramps, $\mathrm{PSD}_{\mathrm{c}}$ increases smoothly exponentially during this part of the sequence (Fig. 2(e)), due in part to the finite thermalization rate.

By shortening the sequence we are asking the algorithm to maximize the cooling speed, which is limited by the lower of the collisional rate $\nu_{\mathrm{c}}$ and the trap vibration frequencies $\nu_{x,y,z}$ [33]. When the gas is still hot, we have $\nu_{\mathrm{c}}\ll\nu_{x,y,z}$, and the algorithm employs Raman cooling to increase the density and collision rate (Fig. 2(d)). However, when $\nu_{\mathrm{c}}$ approaches the lowest trap vibration frequency $\nu_{y}$ near the time $t=225$ ms, the algorithm starts to reduce the Raman rate, and a little later the optical pumping rate, in order to reduce light induced collisions that scale with $\nu_{\mathrm{c}}$, rather than the trap vibration frequency. Subsequently, for times $t>225$ ms, the cooling proceeds near optimally, with the collision rate close to, but a little below, the trap vibration frequencies. Furthermore, as the system approaches condensation near $t=410$ ms, the collision rate is somewhat lowered to reduce light-induced atom loss (Fig. 2(c)).

Another effect limiting the cooling speed is the loading of the atoms from the single horizontal trap, in which the sample is initially prepared, into the crossed dipole trap (see the movie in SM [34]). Initially, the vertical-beam power $P_{z}$ is held low to avoid creating a high-density dimple region which would lead to excess loss during Raman cooling. Later, $P_{z}$ is ramped up to gather atoms from the horizontal trap beam into the overlap region of the cODT in order to increase the collision rate and speed up evaporative cooling. The relatively sudden ramping of the trap power up and then back down visible in Fig. 2a likely involves an optimal-control-like process since the trap compression and relaxation are faster than the axial period of the horizontal trap of $\sim 200$ ms.

The optimization tended to turn off the Raman cooling after five stages because the cloud temperature $T$ was below the effective recoil temperature [25] where Raman cooling, even with optimal parameters, becomes too slow, while leading to trap loss and heating due to light-assisted collisions [35]. The Bayesian optimization recognized this and shut down the Raman cooling at this point, with the atomic gas close to condensation. Subsequently, at higher compression which is primarily achieved by increasing the vertical beam power, the horizontal trap power is reduced and atoms are efficiently evaporated along the direction of gravity in the tilted potential [20] (see the movie in the SM [34]). Note also that once the atoms have been loaded into the crossed-trap region (after $t=350$ ms), the algorithm makes all trap vibration frequencies similar, which provides the fastest overall thermalization, and hence largest cooling speed.

The BEC is fully prepared at the end of the evaporation stages, $575\ \mathrm{ms}$ after the start of the MOT loading. The final cloud contains $3.7\times 10^{3}$ total atoms and is shown in Fig. 3(b). A bimodal fit of the cloud indicates that $2.8\times 10^{3}$ atoms (76 %) are in the BEC. Although the sequence was optimized for speed rather than efficiency, the initial cooling occurs with a logarithmic slope $\gamma=\mathrm{d}(\log\mathrm{PSD}_{\mathrm{c}})/\mathrm{d}(\log N)\approx 16$.

Cost function impact.—The atomic gases produced by sequences optimized for different values of $\alpha$ are presented in Fig. 4, as well as the results when optimizing for total atom number ($N$). Larger values of $\alpha$ result in more atoms, but at higher temperature and lower condensate fraction, while smaller values of $\alpha$ produce purer BECs, but with fewer atoms overall. Setting $\alpha$ to $0.5$ was found to make a reasonable compromise (orange curve in Fig. 4); so that value was used for the final full-sequence optimization which yielded the data presented in Fig. 2.

Outlook.—In conclusion, we have demonstrated that Raman cooling with far detuned optical-pumping light combined with a final evaporation can rapidly produce BECs with a comparatively simple apparatus, even with a standard alkali atom which lacks narrow optical transitions. Bayesian optimization greatly eased the search for a short sequence to BEC, quickly discovering initially unintuitive yet high-performing sequences. Inspection of the parameters chosen by the algorithm reveals several physical strategies, such as adjusting a collision rate close to, but below the trap vibration frequencies to maximize the thermalization and cooling speed while minimizing density-dependent atom loss, non-adiabatic loading into the crossed-trap dimple, and creating a nearly isotropic trap for efficient evaporation. In future applications, faster condensation can likely be achieved by including dynamical tuning of trap size [21], while user intervention may be further reduced by factoring the sequence length into the assigned cost [14]. We anticipate that many other experimental procedures in atomic physics and beyond can be improved by machine learning.