val navController = rememberNavController()

val screen1ViewModel = viewModel<Screen1ViewModel>()
val screen2ViewModel = viewModel<Screen2ViewModel>()
val screen3ViewModel = viewModel<Screen3ViewModel>()

    ...

NavHost(
        navController,
        startDestination = MainScreen.Screen1.route,
        modifier = Modifier.padding(top = 16.dp, bottom = 16.dp)
    ) {

        composable(MainScreen.Screen1.route) { Screen.Screen1(screen1ViewModel) }
        composable(MainScreen.Screen2.route) { Screen.Screen2(screen2ViewModel) }
        composable(MainScreen.Screen3.route) { Screen.Screen3(screen3ViewModel) }
    }
}

interface LazyViewModel<T : ViewModel> {
    @Composable
    fun get(): T
}

class LazyViewModelImpl<T : ViewModel>(
    private val viewModelClass: Class<T>
) : LazyViewModel<T> {
    private var viewModel: T? = null

    @Composable
    override fun get(): T {
        check(Looper.getMainLooper() == Looper.myLooper())

        if (viewModel == null) {
            viewModel = viewModel(viewModelClass)
        }
        return viewModel!!
    }
}

@Composable
inline fun <reified T : ViewModel> lazyViewModel(): LazyViewModel<T> {
    return LazyViewModelImpl(T::class.java)
}

restoreState?