public function step(Parameter $param, Tensor $gradient) : void
    {
        [$velocity, $g2] = $this->cache[$param->id()];

        $velocity = $velocity->multiply($this->momentumDecay)
            ->add($gradient->multiply(1. - $this->momentumDecay));

        $g2 = $g2->multiply($this->rmsDecay)
            ->add($gradient->square()->multiply(1. - $this->rmsDecay));

        $this->cache[$param->id()] = [$velocity, $g2];

        if ($this->t < self::WARM_UP_STEPS) {
            $this->t++;
            
            $velocity = $velocity->divide(1. - $this->momentumDecay ** $this->t);

            $g2 = $g2->divide(1. - $this->rmsDecay ** $this->t);
        }

        $step = $velocity->multiply($this->rate)
            ->divide($g2->sqrt()->clipLower(EPSILON));

        $param->update($step);
    }