Вдохновившись Вашей реализацией решил "позаимствовать" оптимизатор для TF себе в копилку и попутно порефакторить. Вот что нашлось:

Потребление памяти
SGD без момента вообще не использует доп. переменные и как следствие не потребляет дополнительной памяти. С моментом = память x2
Adam если правильно помню хранит 3 слота = память x3
Adan только для dense-реалзиации потребляет x4 памяти, а в Вашей реализации x5 из-за вот этого https://github.com/DenisVorotyntsev/Adan/blob/main/tf_adan/adan.py#L50

Хорошего способа совсем убрать счетчик апдейтов я не нашел, но можно значительно урезать объем этого 5го слота (сценарий когда будут делать срезы внутри канала имхо маловероятен) https://github.com/shkarupa-alex/tfmiss/blob/develop/tfmiss/keras/optimizers/adan.py#L51

Точность sparse-обновлений #1
При расчете каждого обновления Adan использует текущий номер шага в расчете bias_correction_* https://github.com/DenisVorotyntsev/Adan/blob/main/tf_adan/adan.py#L74
При sparse-обновлениях нужно использовать текущий номер апдейта для срезов, а не глобальный номер шага (можно проверить поэлементно пропуская те шаги в которых нет индекса этого элемента)

Чтобы получить точное соответствие dense-ветке пришлось вынести bias_correction_* в каждую из веток и считать их в dense-ветке на основе глобального шага (как в Вашей реализации) а вот в sparse-ветке на основе того самого счетчика обновлений https://github.com/shkarupa-alex/tfmiss/blob/develop/tfmiss/keras/optimizers/adan.py#L148

Точность sparse-обновлений #2
_resource_scatter_update возвращает всю переменную, а не только текущий срез
В Вашей реализации это приводит к обновлению всей переменной каждый раз, а не только к обновлению текущего среза (что во-первых менее производительно, во вторых кажется несет ошибку).
Т.е. если какой-то категориальной переменной в срезе не было она все равно обновится по данным предыдущих итераций.
Пришлось немного переструктурировать код и все sparse-обновления делать после расчетов на текущих срезах

Прочее
Если вот так брать скорость обучения https://github.com/DenisVorotyntsev/Adan/blob/main/tf_adan/adan.py#L59 подозреваю что не будут работать расписания lr
Кажется правильнее брать lr_t который появляется после super()._prepare_local(...)

По графикам не все однозначно (sparse-часть блокнота).
Моя реализация оказывается ближе к оригинальной почти везде кроме 1го графика где ведет себя лучше и стабильнее чем оригинальная и Ваша.

        self._set_hyper("learning_rate", learning_rate)
        self._set_hyper("beta_1", beta_1)
        self._set_hyper("beta_2", beta_2)
        self._set_hyper("beta_3", beta_3)
        self._set_hyper("epsilon", epsilon)
        self._set_hyper("weight_decay", weight_decay)

In init() fucntion, the code used the function _set_hyper(), which was obsoleted .