SM3 Memory Sketch Lite

Problem

Perform one SM3-style update with a fixed coordinate subset.

For each index in subset, replace acc[i] with min(acc[i], g[i] ** 2). Then update w[i] = w[i] - lr * g[i] / (sqrt(acc[i]) + eps) when acc[i] > 0, otherwise w[i] - lr * g[i]. Return [w_new, acc_new].

Hint

Only the selected coordinates refresh their accumulator.

Test cases

case 1

in: [[1,1],[0.5,2],[1,1],0.1,1e-8,[0,1]]

expected: [[0.900000002,0.800000002],[0.25,1]]

case 2

in: [[1,1],[0.5,2],[0.1,4],0.1,1e-8,[0]]

expected: [[0.8418861219915809,0.9000000005],[0.1,4]]

case 3

in: [[0],[3],[0],0.5,1e-8,[0]]

expected: [[-1.5],[0]]

Starter codePython

def sm3_update(w, g, acc, lr, eps, subset):
    # Returns [w_new, acc_new]
    # Your code here
    pass

3 test cases · run your code against them in the editor