|
| 1 | +defmodule Scholar.Optimize.BFGS do |
| 2 | + @moduledoc """ |
| 3 | + BFGS (Broyden-Fletcher-Goldfarb-Shanno) algorithm for multivariate function minimization. |
| 4 | +
|
| 5 | + BFGS is a quasi-Newton optimization method that approximates the inverse Hessian |
| 6 | + matrix using gradient information. It is well-suited for smooth, differentiable |
| 7 | + objective functions and typically converges much faster than derivative-free methods |
| 8 | + like Nelder-Mead. |
| 9 | +
|
| 10 | + ## Algorithm |
| 11 | +
|
| 12 | + At each iteration, the algorithm: |
| 13 | + 1. Computes the gradient using automatic differentiation |
| 14 | + 2. Determines a search direction from the inverse Hessian approximation |
| 15 | + 3. Performs a line search to find an acceptable step length |
| 16 | + 4. Updates the inverse Hessian approximation using the BFGS formula |
| 17 | +
|
| 18 | + ## Convergence |
| 19 | +
|
| 20 | + The algorithm converges when the gradient norm falls below the specified tolerance. |
| 21 | +
|
| 22 | + ## References |
| 23 | +
|
| 24 | + * Nocedal, J. and Wright, S. J. (2006). "Numerical Optimization" |
| 25 | + * Fletcher, R. (1987). "Practical Methods of Optimization" |
| 26 | + """ |
| 27 | + |
| 28 | + import Nx.Defn |
| 29 | + |
| 30 | + @derive {Nx.Container, containers: [:x, :fun, :converged, :iterations, :fun_evals, :grad_evals]} |
| 31 | + defstruct [:x, :fun, :converged, :iterations, :fun_evals, :grad_evals] |
| 32 | + |
| 33 | + @type t :: %__MODULE__{ |
| 34 | + x: Nx.Tensor.t(), |
| 35 | + fun: Nx.Tensor.t(), |
| 36 | + converged: Nx.Tensor.t(), |
| 37 | + iterations: Nx.Tensor.t(), |
| 38 | + fun_evals: Nx.Tensor.t(), |
| 39 | + grad_evals: Nx.Tensor.t() |
| 40 | + } |
| 41 | + |
| 42 | + # Line search parameter (Armijo condition) |
| 43 | + @c1 1.0e-4 |
| 44 | + |
| 45 | + opts = [ |
| 46 | + gtol: [ |
| 47 | + type: {:custom, Scholar.Options, :positive_number, []}, |
| 48 | + default: 1.0e-5, |
| 49 | + doc: """ |
| 50 | + Gradient norm tolerance for convergence. The algorithm stops when |
| 51 | + the infinity norm of the gradient is below this threshold. |
| 52 | + """ |
| 53 | + ], |
| 54 | + maxiter: [ |
| 55 | + type: :pos_integer, |
| 56 | + default: 500, |
| 57 | + doc: "Maximum number of iterations." |
| 58 | + ] |
| 59 | + ] |
| 60 | + |
| 61 | + @opts_schema NimbleOptions.new!(opts) |
| 62 | + |
| 63 | + @doc """ |
| 64 | + Minimizes a multivariate function using the BFGS algorithm. |
| 65 | +
|
| 66 | + BFGS is a gradient-based method that uses automatic differentiation to compute |
| 67 | + gradients and approximates the inverse Hessian matrix for fast convergence. |
| 68 | +
|
| 69 | + ## Arguments |
| 70 | +
|
| 71 | + * `x0` - Initial guess as a 1-D tensor of shape `{n}`. |
| 72 | + * `fun` - The objective function to minimize. Must be a defn-compatible function |
| 73 | + that takes a 1-D tensor of shape `{n}` and returns a scalar tensor. |
| 74 | + * `opts` - Options (see below). |
| 75 | +
|
| 76 | + ## Options |
| 77 | +
|
| 78 | + #{NimbleOptions.docs(@opts_schema)} |
| 79 | +
|
| 80 | + ## Returns |
| 81 | +
|
| 82 | + A `Scholar.Optimize.BFGS` struct with the optimization result: |
| 83 | +
|
| 84 | + * `:x` - The optimal point found (shape `{n}`) |
| 85 | + * `:fun` - The function value at the optimal point |
| 86 | + * `:converged` - Whether the optimization converged (1 if true, 0 if false) |
| 87 | + * `:iterations` - Number of iterations performed |
| 88 | + * `:fun_evals` - Number of function evaluations |
| 89 | + * `:grad_evals` - Number of gradient evaluations |
| 90 | +
|
| 91 | + ## Examples |
| 92 | +
|
| 93 | + iex> # Minimize a simple quadratic: f(x) = x1^2 + x2^2 |
| 94 | + iex> fun = fn x -> Nx.sum(Nx.pow(x, 2)) end |
| 95 | + iex> x0 = Nx.tensor([1.0, 2.0]) |
| 96 | + iex> result = Scholar.Optimize.BFGS.minimize(x0, fun) |
| 97 | + iex> Nx.to_number(result.converged) |
| 98 | + 1 |
| 99 | + iex> Nx.all_close(result.x, Nx.tensor([0.0, 0.0]), atol: 1.0e-4) |> Nx.to_number() |
| 100 | + 1 |
| 101 | +
|
| 102 | + For higher precision, use f64 tensors: |
| 103 | +
|
| 104 | + iex> fun = fn x -> Nx.sum(Nx.pow(x, 2)) end |
| 105 | + iex> x0 = Nx.tensor([1.0, 2.0], type: :f64) |
| 106 | + iex> result = Scholar.Optimize.BFGS.minimize(x0, fun, gtol: 1.0e-10) |
| 107 | + iex> Nx.to_number(result.converged) |
| 108 | + 1 |
| 109 | + iex> Nx.all_close(result.x, Nx.tensor([0.0, 0.0]), atol: 1.0e-8) |> Nx.to_number() |
| 110 | + 1 |
| 111 | +
|
| 112 | + Minimizing the Rosenbrock function: |
| 113 | +
|
| 114 | + iex> # Rosenbrock: f(x,y) = (1-x)^2 + 100*(y-x^2)^2, minimum at (1, 1) |
| 115 | + iex> rosenbrock = fn x -> |
| 116 | + ...> x0 = x[0] |
| 117 | + ...> x1 = x[1] |
| 118 | + ...> term1 = Nx.pow(Nx.subtract(1, x0), 2) |
| 119 | + ...> term2 = Nx.multiply(100, Nx.pow(Nx.subtract(x1, Nx.pow(x0, 2)), 2)) |
| 120 | + ...> Nx.add(term1, term2) |
| 121 | + ...> end |
| 122 | + iex> x0 = Nx.tensor([0.0, 0.0], type: :f64) |
| 123 | + iex> result = Scholar.Optimize.BFGS.minimize(x0, rosenbrock, gtol: 1.0e-8, maxiter: 1000) |
| 124 | + iex> Nx.to_number(result.converged) |
| 125 | + 1 |
| 126 | + iex> Nx.all_close(result.x, Nx.tensor([1.0, 1.0]), atol: 1.0e-4) |> Nx.to_number() |
| 127 | + 1 |
| 128 | + """ |
| 129 | + deftransform minimize(x0, fun, opts \\ []) do |
| 130 | + opts = NimbleOptions.validate!(opts, @opts_schema) |
| 131 | + minimize_n(x0, fun, opts[:gtol], opts[:maxiter]) |
| 132 | + end |
| 133 | + |
| 134 | + defnp minimize_n(x0, fun, gtol, maxiter) do |
| 135 | + x0 = Nx.flatten(x0) |
| 136 | + {n} = Nx.shape(x0) |
| 137 | + |
| 138 | + # Initial function value and gradient |
| 139 | + {f0, g0} = value_and_grad(x0, fun) |
| 140 | + |
| 141 | + # Initialize inverse Hessian as identity matrix |
| 142 | + h_inv = Nx.eye(n, type: Nx.type(x0)) |
| 143 | + |
| 144 | + # Initial state |
| 145 | + initial_state = %{ |
| 146 | + x: x0, |
| 147 | + f: f0, |
| 148 | + g: g0, |
| 149 | + h_inv: h_inv, |
| 150 | + iter: Nx.u32(0), |
| 151 | + f_evals: Nx.u32(1), |
| 152 | + g_evals: Nx.u32(1) |
| 153 | + } |
| 154 | + |
| 155 | + # Main optimization loop |
| 156 | + {final_state, _} = |
| 157 | + while {state = initial_state, {gtol, maxiter}}, |
| 158 | + not converged?(state, gtol) and state.iter < maxiter do |
| 159 | + new_state = bfgs_step(state, fun) |
| 160 | + {new_state, {gtol, maxiter}} |
| 161 | + end |
| 162 | + |
| 163 | + # Check convergence |
| 164 | + converged = converged?(final_state, gtol) |
| 165 | + |
| 166 | + %__MODULE__{ |
| 167 | + x: final_state.x, |
| 168 | + fun: final_state.f, |
| 169 | + converged: converged, |
| 170 | + iterations: final_state.iter, |
| 171 | + fun_evals: final_state.f_evals, |
| 172 | + grad_evals: final_state.g_evals |
| 173 | + } |
| 174 | + end |
| 175 | + |
| 176 | + # Check convergence: infinity norm of gradient < gtol |
| 177 | + defnp converged?(state, gtol) do |
| 178 | + grad_norm = Nx.reduce_max(Nx.abs(state.g)) |
| 179 | + Nx.less(grad_norm, gtol) |
| 180 | + end |
| 181 | + |
| 182 | + # Perform one BFGS iteration |
| 183 | + defnp bfgs_step(state, fun) do |
| 184 | + x = state.x |
| 185 | + f = state.f |
| 186 | + g = state.g |
| 187 | + h_inv = state.h_inv |
| 188 | + |
| 189 | + # Compute search direction: p = -H_inv * g |
| 190 | + p = Nx.negate(Nx.dot(h_inv, g)) |
| 191 | + |
| 192 | + # Line search to find step length alpha |
| 193 | + {alpha, f_new, g_new, ls_f_evals, ls_g_evals} = |
| 194 | + line_search(x, f, g, p, fun) |
| 195 | + |
| 196 | + # Compute step and gradient change |
| 197 | + s = Nx.multiply(alpha, p) |
| 198 | + x_new = Nx.add(x, s) |
| 199 | + y = Nx.subtract(g_new, g) |
| 200 | + |
| 201 | + # Update inverse Hessian using BFGS formula |
| 202 | + h_inv_new = update_inverse_hessian(h_inv, s, y) |
| 203 | + |
| 204 | + %{ |
| 205 | + state |
| 206 | + | x: x_new, |
| 207 | + f: f_new, |
| 208 | + g: g_new, |
| 209 | + h_inv: h_inv_new, |
| 210 | + iter: Nx.add(state.iter, 1), |
| 211 | + f_evals: Nx.add(state.f_evals, ls_f_evals), |
| 212 | + g_evals: Nx.add(state.g_evals, ls_g_evals) |
| 213 | + } |
| 214 | + end |
| 215 | + |
| 216 | + # Backtracking line search with Armijo condition |
| 217 | + defnp line_search(x, f, g, p, fun) do |
| 218 | + # Directional derivative at alpha=0 |
| 219 | + slope = Nx.dot(g, p) |
| 220 | + |
| 221 | + # Try step sizes: 1, 0.5, 0.25, 0.125, ... |
| 222 | + # Unrolled at compile time for fusion. |
| 223 | + {alpha, _} = |
| 224 | + while {alpha = Nx.tensor(1.0, type: Nx.type(x)), {x, f, p, slope}}, |
| 225 | + _i <- 0..9//1, |
| 226 | + unroll: true do |
| 227 | + new_alpha = backtrack_step(x, f, p, slope, fun, alpha) |
| 228 | + {new_alpha, {x, f, p, slope}} |
| 229 | + end |
| 230 | + |
| 231 | + # Evaluate function and gradient at final alpha |
| 232 | + x_final = Nx.add(x, Nx.multiply(alpha, p)) |
| 233 | + {f_final, g_final} = value_and_grad(x_final, fun) |
| 234 | + |
| 235 | + {alpha, f_final, g_final, Nx.u32(11), Nx.u32(1)} |
| 236 | + end |
| 237 | + |
| 238 | + # Single backtracking step |
| 239 | + defnp backtrack_step(x, f, p, slope, fun, alpha) do |
| 240 | + x_trial = Nx.add(x, Nx.multiply(alpha, p)) |
| 241 | + f_trial = fun.(x_trial) |
| 242 | + |
| 243 | + # Check Armijo condition |
| 244 | + armijo_ok = |
| 245 | + Nx.less_equal(f_trial, Nx.add(f, Nx.multiply(@c1, Nx.multiply(alpha, slope)))) |
| 246 | + |
| 247 | + # If satisfied, keep alpha; otherwise halve it |
| 248 | + Nx.select(armijo_ok, alpha, Nx.multiply(0.5, alpha)) |
| 249 | + end |
| 250 | + |
| 251 | + # BFGS inverse Hessian update |
| 252 | + # H_new = (I - rho*s*y') * H * (I - rho*y*s') + rho*s*s' |
| 253 | + defnp update_inverse_hessian(h_inv, s, y) do |
| 254 | + {n} = Nx.shape(s) |
| 255 | + |
| 256 | + # rho = 1 / (y' * s) |
| 257 | + ys = Nx.dot(y, s) |
| 258 | + |
| 259 | + # Skip update if ys is too small (would cause numerical issues) |
| 260 | + skip_update = Nx.less(Nx.abs(ys), 1.0e-10) |
| 261 | + |
| 262 | + rho = Nx.select(skip_update, 0.0, Nx.divide(1.0, ys)) |
| 263 | + |
| 264 | + # Compute update terms |
| 265 | + # I - rho*s*y' and I - rho*y*s' |
| 266 | + eye = Nx.eye(n, type: Nx.type(s)) |
| 267 | + |
| 268 | + # s and y as column vectors for outer products |
| 269 | + s_col = Nx.reshape(s, {n, 1}) |
| 270 | + y_col = Nx.reshape(y, {n, 1}) |
| 271 | + s_row = Nx.reshape(s, {1, n}) |
| 272 | + y_row = Nx.reshape(y, {1, n}) |
| 273 | + |
| 274 | + # (I - rho*s*y') |
| 275 | + left = Nx.subtract(eye, Nx.multiply(rho, Nx.dot(s_col, y_row))) |
| 276 | + |
| 277 | + # (I - rho*y*s') |
| 278 | + right = Nx.subtract(eye, Nx.multiply(rho, Nx.dot(y_col, s_row))) |
| 279 | + |
| 280 | + # rho*s*s' |
| 281 | + ss_term = Nx.multiply(rho, Nx.dot(s_col, s_row)) |
| 282 | + |
| 283 | + # H_new = left * H * right + ss_term |
| 284 | + h_new = Nx.add(Nx.dot(Nx.dot(left, h_inv), right), ss_term) |
| 285 | + |
| 286 | + # Use old H if update skipped |
| 287 | + Nx.select(skip_update, h_inv, h_new) |
| 288 | + end |
| 289 | +end |
0 commit comments