Abstract: The physical properties of the outer layers of stellar atmospheres (temperature, velocity and/or magnetic field) can be inferred by inverting the radiative transfer forward problem. The main obstacle is that the model required to synthesize the strong lines that sample the stellar chromospheres is extremely time consuming, which makes the solution of the inverse problem not very practical. Here we leverage graph networks to predict the population number density of the atom energy levels simply from the temperature and optical depth stratification. We demonstrate that a speedup of a factor 10$^3$ can be obtained with a negligible impact on precision. This opens up the possibility of large-scale synthesis in three-dimensional models and routine inversion of observations to infer the 3D properties of the solar and stellar chromospheres.