src/modules/optim/lion.cc¶

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "tenzr/modules/optim/lion.h"

#include <cmath>
#include <vector>

#include "tenzr/common/error.h"

namespace tenzr::optim {

  Lion::Lion(double lr, double beta1, double beta2, double weight_decay, ScheduleFn schedule)
      : lr_(lr), beta1_(beta1), beta2_(beta2), weight_decay_(weight_decay), schedule_(std::move(schedule)) {
    TENZR_CHECK(lr >= 0.0, "Learning rate must be non-negative");
    TENZR_CHECK(beta1 >= 0.0 && beta1 < 1.0, "Beta1 must be in [0, 1)");
    TENZR_CHECK(beta2 >= 0.0 && beta2 < 1.0, "Beta2 must be in [0, 1)");
  }

  OptState Lion::Init(const std::unordered_map<std::string, Tensor>& params) {
    OptState state;
    state.step = 0;
    for (const auto& [name, p] : params) {
      state.param_states[name]["m"] = Tensor::Zeros(p.shape(), p.Dtype());
    }
    return state;
  }

  std::unordered_map<std::string, Tensor> Lion::Update(const std::unordered_map<std::string, Tensor>& grads, OptState& state,
                                                       const std::unordered_map<std::string, Tensor>& params) {
    double current_lr = schedule_ ? schedule_(state.step) : lr_;
    double lr_d = current_lr;
    double b1_d = beta1_;
    double b2_d = beta2_;
    double wd_d = weight_decay_;

    std::unordered_map<std::string, Tensor> updates;

    for (const auto& [name, g] : grads) {
      auto p_it = params.find(name);
      TENZR_CHECK(p_it != params.end(), "Gradient key not found in parameters: " + name);
      const auto& p = p_it->second;

      // A lazy gradient is consumed here (to_vector materializes it).
      if (!p.HasData()) {
        updates[name] = g;
        continue;
      }

      auto g_raw = g.ToVectorAs<double>();
      auto p_raw = p.ToVectorAs<double>();
      size_t n = g_raw.size();

      auto& m_tensor = state.param_states[name]["m"];
      auto m_raw = m_tensor.ToVectorAs<double>();
      std::vector<double> u_raw(n);

      for (size_t i = 0; i < n; ++i) {
        double grad_val = g_raw[i];

        // Interpolate for update direction: c_t = beta1 * m_{t-1} + (1 - beta1) * g_t
        double c = b1_d * m_raw[i] + (1.0 - b1_d) * grad_val;

        // Sign of c
        double s = (c > 0.0) ? 1.0 : ((c < 0.0) ? -1.0 : 0.0);

        double step_val = -lr_d * s;
        if (wd_d > 0.0) {
          step_val -= lr_d * wd_d * p_raw[i];
        }

        u_raw[i] = step_val;

        // Update momentum: m_t = beta2 * m_{t-1} + (1 - beta2) * g_t
        m_raw[i] = b2_d * m_raw[i] + (1.0 - b2_d) * grad_val;
      }

      m_tensor.BindData(Tensor::FromVector<double>(m_raw, m_tensor.shape()).Astype(m_tensor.Dtype()).Data());

      updates[name] = Tensor::FromVector<double>(u_raw, p.shape()).Astype(p.Dtype());
    }

    state.step++;
    return updates;
  }

}  // namespace tenzr::optim