| #------------------------------------------------------------- |
| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you 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 |
| # |
| # http://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. |
| # |
| #------------------------------------------------------------- |
| |
| /* |
| * Shampoo optimizer. |
| * |
| * Implementation of the Shampoo optimizer as described in: |
| * |
| * Gupta et al., "Shampoo: Preconditioned Stochastic Tensor Optimization" |
| * https://arxiv.org/abs/1802.09568 |
| * |
| * Shampoo is a second-order optimization method that preconditions |
| * gradients using estimates of the row- and column-wise covariance |
| * of the gradients. Compared to first-order optimizers (SGD, Adam), |
| * Shampoo can converge faster but is significantly more memory-intensive. |
| * |
| * This implementation supports: |
| * - Full-matrix Shampoo (exact preconditioning) |
| * - Diagonal Shampoo (memory-efficient approximation) |
| * |
| * The choice between the two modes is determined by the shape of X |
| * and the preconditioner initialization. |
| */ |
| |
| update = function(matrix[double] X, matrix[double] dX, double lr, |
| matrix[double] preconL, matrix[double] preconR, boolean useDiag) |
| return(matrix[double] X, matrix[double] preconL, matrix[double] preconR){ |
| /* |
| * Performs one optimization step using the Shampoo update rule. |
| * |
| * |
| * Inputs: |
| * - X: Parameter matrix to be updated (n × m) |
| * - dX: Gradient of the loss w.r.t. X (n × m) |
| * - lr: Learning rate. |
| * - preconL: Left (row) preconditioner |
| * - Full: (n × n) |
| * - Diagonal: (n × 1) |
| * - preconR: Right (column) preconditioner |
| * - Full: (m × m) |
| * - Diagonal: (1 × m) |
| * - useDiag: Boolean flag indicating whether diagonal Shampoo is used |
| * |
| * Outputs: |
| * - X: Updated parameter matrix (n × m) |
| * - preconL: Updated left preconditioner |
| * - Full: (n × n) |
| * - Diagonal: (n × 1) |
| * - preconR: Updated right preconditioner |
| * - Full: (m × m) |
| * - Diagonal: (1 × m) |
| */ |
| |
| # Full-matrix Shampoo: |
| # Only used if both dimensions are small enough |
| if(useDiag==FALSE){ |
| |
| preconL = preconL + dX %*% t(dX) |
| preconR = preconR + t(dX) %*% dX |
| |
| [LEigenvalue, LEigenvector] = eigen(preconL) |
| preconLInvPowerRoot = LEigenvector %*% diag(LEigenvalue^(-0.25)) %*% t(LEigenvector) |
| |
| [REigenvalue, REigenvector] = eigen(preconR) |
| preconRInvPowerRoot = REigenvector %*% diag(REigenvalue^(-0.25)) %*% t(REigenvector) |
| |
| X = X - lr * preconLInvPowerRoot %*% dX %*% preconRInvPowerRoot |
| |
| # Diagonal Shampoo: |
| # Memory-efficient approximation for large parameter matrices |
| } else{ |
| n = nrow(dX) |
| m = ncol(dX) |
| |
| preconL = preconL + rowSums(dX^2) |
| preconR = preconR + colSums(dX^2) |
| |
| preconLScale = preconL^(-0.25) |
| preconRScale = preconR^(-0.25) |
| |
| preconLMatrix = preconLScale %*% matrix(1, rows=1, cols=m) |
| preconRMatrix = matrix(1, rows=n, cols=1) %*% preconRScale |
| |
| scaledGrad = dX * preconLMatrix; |
| scaledGrad = scaledGrad * preconRMatrix; |
| |
| X = X - lr * scaledGrad; |
| |
| } |
| } |
| |
| init = function(matrix[double] X, double epsilon, int useDiagThreshold) |
| return (matrix[double] preconL, matrix[double] preconR, boolean useDiag) { |
| /* |
| * Initializes the Shampoo preconditioners for a given parameter matrix. |
| * |
| * Depending on the size of X, this function initializes either: |
| * - Full identity matrices (exact Shampoo), or |
| * - Diagonal vectors (approximate Shampoo) |
| * |
| * This threshold is crucial to avoid excessive memory usage, |
| * as full Shampoo requires O(n^2 + m^2) memory per parameter matrix. |
| * |
| * Inputs: |
| * - X: Parameter matrix to be optimized (n, m) |
| * - epsilon: Numerical stability constant |
| * - useDiagThreshold: Dimension threshold above which diagonal |
| * preconditioning is used |
| * |
| * Outputs: |
| * - preconL: Initial left preconditioner |
| * - Full: (n × n) identity scaled by epsilon |
| * - Diagonal: (n × 1) filled with epsilon |
| * - preconR: Initial right preconditioner |
| * - Full: (m × m) identity scaled by epsilon |
| * - Diagonal: (1 × m) filled with epsilon |
| * - useDiag: Boolean flag indicating whether diagonal Shampoo is used |
| */ |
| |
| # Use diagonal Shampoo if parameter matrix is too large |
| if((nrow(X) > useDiagThreshold) | (ncol(X) > useDiagThreshold)){ |
| preconL = matrix(epsilon, rows=nrow(X), cols=1); |
| preconR = matrix(epsilon, rows=1, cols=ncol(X)); |
| useDiag = TRUE |
| |
| # Use full Shampoo if parameter matrix is small enough |
| } else { |
| preconL = matrix(0, rows=nrow(X), cols=nrow(X)); |
| index = 1; |
| while (index <= nrow(X)){ |
| preconL[index, index] = epsilon * 1 |
| index = index + 1 |
| } |
| preconR = matrix(0, rows=ncol(X), cols=ncol(X)); |
| index = 1; |
| while (index <= ncol(X)){ |
| preconR[index, index] = epsilon * 1 |
| index = index + 1 |
| } |
| useDiag = FALSE |
| } |
| } |
| |
| update_momentum = function(matrix[double] X, matrix[double] dX, double lr, |
| matrix[double] preconL, matrix[double] preconR, |
| matrix[double] momentum, boolean useDiag) |
| return(matrix[double] X, matrix[double] preconL, matrix[double] preconR, |
| matrix[double] momentum){ |
| /* |
| * Performs one optimization step using the Shampoo update rule, while using momentum. |
| * |
| * |
| * Inputs: |
| * - X: Parameter matrix to be updated (n × m) |
| * - dX: Gradient of the loss w.r.t. X (n × m) |
| * - lr: Learning rate. |
| * - preconL: Left (row) preconditioner |
| * - Full: (n × n) |
| * - Diagonal: (n × 1) |
| * - preconR: Right (column) preconditioner |
| * - Full: (m × m) |
| * - Diagonal: (1 × m) |
| * - momentum: momentum (n × m) |
| * - useDiag: Boolean flag indicating whether diagonal Shampoo is used |
| * |
| * Outputs: |
| * - X: Updated parameter matrix (n × m) |
| * - preconL: Updated left preconditioner |
| * - Full: (n × n) |
| * - Diagonal: (n × 1) |
| * - preconR: Updated right preconditioner |
| * - Full: (m × m) |
| * - Diagonal: (1 × m) |
| * - momentum: Updated momentum (n × m) |
| */ |
| |
| # calculating the updated momentum |
| momentum = 0.9 * momentum + (0.1)*dX |
| |
| # Full-matrix Shampoo: |
| # Only used if both dimensions are small enough |
| if(useDiag==FALSE){ |
| |
| preconL = preconL + dX %*% t(dX) |
| preconR = preconR + t(dX) %*% dX |
| |
| [LEigenvalue, LEigenvector] = eigen(preconL) |
| preconLInvPowerRoot = LEigenvector %*% diag(LEigenvalue^(-0.25)) %*% t(LEigenvector) |
| |
| [REigenvalue, REigenvector] = eigen(preconR) |
| preconRInvPowerRoot = REigenvector %*% diag(REigenvalue^(-0.25)) %*% t(REigenvector) |
| |
| X = X - lr * preconLInvPowerRoot %*% momentum %*% preconRInvPowerRoot |
| |
| # Diagonal Shampoo: |
| # Memory-efficient approximation for large parameter matrices |
| } else{ |
| n = nrow(dX) |
| m = ncol(dX) |
| |
| preconL = preconL + rowSums(dX ^ 2) |
| preconR = preconR + colSums(dX ^ 2) |
| |
| preconLScale = preconL^(-0.25) |
| preconRScale = preconR^(-0.25) |
| |
| preconLMatrix = preconLScale %*% matrix(1, rows=1, cols=m) |
| preconRMatrix = matrix(1, rows=n, cols=1) %*% preconRScale |
| |
| scaledGrad = momentum * preconLMatrix |
| scaledGrad = scaledGrad * preconRMatrix |
| |
| X = X - lr * scaledGrad |
| } |
| } |
| |
| init_momentum = function(matrix[double] X, double epsilon, int useDiagThreshold) |
| return (matrix[double] preconL, matrix[double] preconR, |
| matrix[double] momentum, boolean useDiag) { |
| /* |
| * Initializes the Shampoo preconditioners and momentum for a given parameter matrix. |
| * |
| * Depending on the size of X, this function initializes either: |
| * - Full identity matrices (exact Shampoo), or |
| * - Diagonal vectors (approximate Shampoo) |
| * |
| * This threshold is crucial to avoid excessive memory usage, |
| * as full Shampoo requires O(n² + m²) memory per parameter matrix. |
| * |
| * Inputs: |
| * - X: Parameter matrix to be optimized (n, m) |
| * - epsilon: Numerical stability constant |
| * - useDiagThreshold: Dimension threshold above which diagonal |
| * preconditioning is used |
| * |
| * Outputs: |
| * - preconL: Initial left preconditioner |
| * - Full: (n × n) identity scaled by epsilon |
| * - Diagonal: (n × 1) filled with epsilon |
| * - preconR: Initial right preconditioner |
| * - Full: (m × m) identity scaled by epsilon |
| * - Diagonal: (1 × m) filled with epsilon |
| * - momentum: Initial momentum (n × m), initialized to zeros |
| * - useDiag: Boolean flag indicating whether diagonal Shampoo is used |
| */ |
| |
| # Use diagonal Shampoo if parameter matrix is too large |
| if((nrow(X) > useDiagThreshold) | (ncol(X) > useDiagThreshold)){ |
| preconL = matrix(epsilon, rows=nrow(X), cols=1); |
| preconR = matrix(epsilon, rows=1, cols=ncol(X)); |
| useDiag = TRUE |
| |
| # Use full Shampoo if parameter matrix is small enough |
| } else { |
| preconL = matrix(0, rows=nrow(X), cols=nrow(X)); |
| index = 1; |
| while (index <= nrow(X)){ |
| preconL[index, index] = epsilon * 1 |
| index = index + 1 |
| } |
| preconR = matrix(0, rows=ncol(X), cols=ncol(X)); |
| index = 1; |
| while (index <= ncol(X)){ |
| preconR[index, index] = epsilon * 1 |
| index = index + 1 |
| } |
| useDiag = FALSE |
| } |
| momentum = X * 0 |
| } |
| |
| update_heuristic = function(matrix[double] X, matrix[double] dX, double lr, |
| matrix[double] preconL, matrix[double] preconR, matrix[double] momentum, |
| int stepCounter, int rootEvery, int preconEvery, matrix[double] bufferL, |
| matrix[double] bufferR, matrix[double] preconLInvPowerRoot, |
| matrix[double] preconRInvPowerRoot, boolean useDiag) |
| return (matrix[double] X, matrix[double] preconL, matrix[double] preconR, |
| matrix[double] momentum, int stepCounter, matrix[double] bufferL, |
| matrix[double] bufferR, matrix[double] preconLInvPowerRoot, |
| matrix[double] preconRInvPowerRoot){ |
| /* |
| * Performs one optimization step using the Shampoo update rule, while using momentum |
| * and a heuristic for runtime improvements. |
| * |
| * |
| * Inputs: |
| * - X: Parameter matrix to be updated (n × m) |
| * - dX: Gradient of the loss w.r.t. X (n × m) |
| * - lr: Learning rate. |
| * - preconL: Left (row) preconditioner |
| * - Full: (n × n) |
| * - Diagonal: (n × 1) |
| * - preconR: Right (column) preconditioner |
| * - Full: (m × m) |
| * - Diagonal: (1 × m) |
| * - momentum: momentum (n × m) |
| * - stepCounter: Step counter (int), incremented each call |
| * - rootEvery: Frequency for recomputing inverse roots (int) |
| * - preconEvery: Frequency for applying buffered updates to preconditioners (int) |
| * - bufferL: Buffer accumulating left curvature updates |
| * - Full: (n × n) |
| * - Diagonal: (n × 1) |
| * - bufferR: Buffer accumulating right curvature updates |
| * - Full: (m × m) |
| * - Diagonal: (1 × m) |
| * - preconLInvPowerRoot: Cached preconL^{-1/4} |
| * - Full: (n × n) |
| * - Diagonal: (n × 1) |
| * - preconRInvPowerRoot: Cached preconR^{-1/4} |
| * - Full: (m × m) |
| * - Diagonal: (1 × m) |
| * - useDiag: Boolean flag indicating whether diagonal Shampoo is used |
| * |
| * Outputs: |
| * - X: Updated parameter matrix (n × m) |
| * - preconL: Updated left preconditioner |
| * - Full: (n × n) |
| * - Diagonal: (n × 1) |
| * - preconR: Updated right preconditioner |
| * - Full: (m × m) |
| * - Diagonal: (1 × m) |
| * - momentum: momentum (n × m) |
| * - stepCounter: Updated step counter (int) |
| * - bufferL: Updated bufferL (reset to 0 when applied) |
| * - Full: (n × n) |
| * - Diagonal: (n × 1) |
| * - bufferR: Updated bufferR (reset to 0 when applied) |
| * - Full: (m × m) |
| * - Diagonal: (1 × m) |
| * - preconLInvPowerRoot: Updated cached inverse root (when recomputed) |
| * - Full: (n × n) |
| * - Diagonal: (n × 1) |
| * - preconRInvPowerRoot: Updated cached inverse root (when recomputed) |
| * - Full: (m × m) |
| * - Diagonal: (1 × m) |
| */ |
| |
| # calculating the updated momentum |
| momentum = 0.9 * momentum + (0.1)*dX |
| |
| # Full-matrix Shampoo: |
| # Only used if both dimensions are small enough |
| if(useDiag==FALSE){ |
| bufferL = bufferL + (dX %*% t(dX)) |
| bufferR = bufferR + (t(dX) %*% dX) |
| |
| if ((stepCounter > 0) & (stepCounter %% preconEvery == 0)){ |
| preconL = preconL + bufferL |
| preconR = preconR + bufferR |
| bufferL = bufferL * 0 |
| bufferR = bufferR * 0 |
| } |
| |
| |
| if ((stepCounter > 0) & (stepCounter %% rootEvery == 0)){ |
| [LEigenvalue, LEigenvector] = eigen(preconL) |
| preconLInvPowerRoot = LEigenvector %*% diag(LEigenvalue^(-0.25)) %*% t(LEigenvector) |
| |
| [REigenvalue, REigenvector] = eigen(preconR) |
| preconRInvPowerRoot = REigenvector %*% diag(REigenvalue^(-0.25)) %*% t(REigenvector) |
| } |
| |
| X = X - lr * preconLInvPowerRoot %*% momentum %*% preconRInvPowerRoot |
| |
| stepCounter = stepCounter + 1 |
| |
| # Diagonal Shampoo: |
| # Memory-efficient approximation for large parameter matrices |
| } else{ |
| n = nrow(dX) |
| m = ncol(dX) |
| |
| bufferL = bufferL + rowSums(dX ^ 2) |
| bufferR = bufferR + colSums(dX ^ 2) |
| |
| if ((stepCounter > 0) & (stepCounter %% preconEvery == 0)){ |
| preconL = preconL + bufferL |
| preconR = preconR + bufferR |
| bufferL = bufferL * 0 |
| bufferR = bufferR * 0 |
| } |
| |
| if ((stepCounter > 0) & (stepCounter %% rootEvery == 0)){ |
| preconLInvPowerRoot = (preconL^(-0.25)) |
| preconRInvPowerRoot = (preconR^(-0.25)) |
| } |
| preconLMatrix = preconLInvPowerRoot %*% matrix(1, rows=1, cols=m) |
| preconRMatrix = matrix(1, rows=n, cols=1) %*% preconRInvPowerRoot |
| |
| scaledGrad = momentum * preconLMatrix |
| scaledGrad = scaledGrad * preconRMatrix |
| |
| X = X - lr * scaledGrad |
| stepCounter = stepCounter + 1 |
| } |
| } |
| |
| |
| |
| init_heuristic = function(matrix[double] X, double epsilon, int useDiagThreshold) |
| return (matrix[double] preconL, matrix[double] preconR, int stepCounter, |
| matrix[double] bufferL, matrix[double] bufferR, matrix[double] momentum, |
| matrix[double] preconLInvPowerRoot, matrix[double] preconRInvPowerRoot, |
| boolean useDiag) { |
| /* |
| * Initializes Shampoo preconditioners, buffers, cached inverse roots, |
| * and momentum for the heuristic variant. |
| * |
| * Depending on the size of X, this function initializes either: |
| * - Full identity matrices (exact Shampoo), or |
| * - Diagonal vectors (approximate Shampoo) |
| * |
| * This threshold is crucial to avoid excessive memory usage, |
| * as full Shampoo requires O(n^2 + m^2) memory per parameter matrix. |
| * |
| * Inputs: |
| * - X: Parameter matrix to be optimized (n, m) |
| * - epsilon: Numerical stability constant |
| * - useDiagThreshold: Dimension threshold above which diagonal |
| * preconditioning is used |
| * |
| * Outputs: |
| * - preconL: Initial left preconditioner (n × n) or (n × 1) |
| * - preconR: Initial right preconditioner (m × m) or (1 × m) |
| * - stepCounter: Initialized to 0 |
| * - bufferL: Initialized to zeros, same shape as preconL |
| * - bufferR: Initialized to zeros, same shape as preconR |
| * - momentum: Initialized to zeros, same shape as X (n × m) |
| * - preconLInvPowerRoot: Cached inverse fourth root of preconL |
| * - Full: initialized to epsilon^{-1/4} * I (n × n) |
| * - Diagonal: initialized to preconL^{-1/4} (n × 1) |
| * - preconRInvPowerRoot: Cached inverse fourth root of preconR |
| * - Full: initialized to epsilon^{-1/4} * I (m × m) |
| * - Diagonal: initialized to preconR^{-1/4} (1 × m) |
| * - useDiag: Boolean flag indicating whether diagonal Shampoo is used |
| */ |
| |
| # Use diagonal Shampoo if parameter matrix is too large |
| if((nrow(X) > useDiagThreshold) | (ncol(X) > useDiagThreshold)){ |
| preconL = matrix(epsilon, rows=nrow(X), cols=1); |
| preconR = matrix(epsilon, rows=1, cols=ncol(X)); |
| preconLInvPowerRoot = preconL^(-0.25) |
| preconRInvPowerRoot = preconR^(-0.25) |
| useDiag = TRUE |
| |
| # Use full Shampoo if parameter matrix is small enough |
| } else { |
| preconL = matrix(0, rows=nrow(X), cols=nrow(X)); |
| index = 1; |
| while (index <= nrow(X)){ |
| preconL[index, index] = epsilon * 1 |
| index = index + 1 |
| } |
| preconR = matrix(0, rows=ncol(X), cols=ncol(X)); |
| index = 1; |
| while (index <= ncol(X)){ |
| preconR[index, index] = epsilon * 1 |
| index = index + 1 |
| } |
| |
| preconLInvPowerRoot = preconL |
| i = 1 |
| while(i <= nrow(preconLInvPowerRoot)) { |
| preconLInvPowerRoot[i,i] = epsilon^(-0.25) |
| i = i + 1 |
| } |
| |
| preconRInvPowerRoot = preconR |
| j = 1 |
| while(j <= nrow(preconRInvPowerRoot)) { |
| preconRInvPowerRoot[j,j] = epsilon^(-0.25) |
| j = j + 1 |
| } |
| |
| useDiag = FALSE |
| } |
| bufferR = preconR * 0 |
| bufferL = preconL * 0 |
| stepCounter = 0 |
| momentum = X * 0 |
| } |
