Computing logits inside train loop - no learning, computing logits inside loss function - works.

[Krupskis]

I was playing with a simple MLP example, in the below example I compute logits inside the train loop - the loss changes, but fluctuates around 2.3.

def evaluate_loss(X, Y):
    return mx.mean(nn.losses.cross_entropy(X, Y))


def eval_fn(X, Y):
    return mx.mean(mx.argmax(X, axis=1) == Y)


if __name__ == "__main__":
    batch_size = 256
    X_train, Y_train, X_test, Y_test = map(mx.array, fetch_mnist())

    model = TinyBobNet()
    optimizer = optim.SGD(learning_rate=1e-1)
    mx.eval(model.parameters())
    loss_and_grad_fn = nn.value_and_grad(model, evaluate_loss)
    count = 0
    for _ in range(500):
        ix = mx.random.randint(0, X_train.shape[0], [batch_size])

        Xb = X_train[ix]
        Yb = Y_train[ix]

        logits = model(Xb)

        loss, grads = loss_and_grad_fn(logits, Yb)
        optimizer.update(model, grads)
        mx.eval(model.parameters(), optimizer.state)
        print(loss)

    print(eval_fn(model(X_test), Y_test))

However, passing model and computing logits inside evaluate_loss fixes it:

def evaluate_loss(model, X, Y):
    return mx.mean(nn.losses.cross_entropy(model(X), Y))


def eval_fn(X, Y):
    return mx.mean(mx.argmax(X, axis=1) == Y)


if __name__ == "__main__":
    batch_size = 256
    X_train, Y_train, X_test, Y_test = map(mx.array, fetch_mnist())

    model = TinyBobNet()
    optimizer = optim.SGD(learning_rate=1e-1)
    mx.eval(model.parameters())
    loss_and_grad_fn = nn.value_and_grad(model, evaluate_loss)
    count = 0
    for _ in range(500):
        ix = mx.random.randint(0, X_train.shape[0], [batch_size])

        Xb = X_train[ix]
        Yb = Y_train[ix]

        # logits = model(Xb)

        loss, grads = loss_and_grad_fn(model, Xb, Yb)
        optimizer.update(model, grads)
        mx.eval(model.parameters(), optimizer.state)
        print(loss)

    print(eval_fn(model(X_test), Y_test))

Could somebody help me understand this behavior?

[awni]

First make sure you understand the grad transform. It takes the derivative of the output w.r.t. to the input. For example y = mx.grad(fn)(x), then y will be the derivatie of fn evaluated at x. We can extend the idea to a dictionary of inputs y = mx.grad(fn)(dictionary_of_inputs) and here y will be the gradient of fn evaluated at dictionary_of_inputs.

Once that makes sense, checkout the docs on nn.value_and_grad which is purely a convenience function to make it easy to get the gradient of the neural net parameters.

It's equivalent to:

def f(params, x):
    model.update(params)  # <---- Necessary to make the model use the passed parameters
    return model(x)

mx.value_and_grad(f)(model.trainable_parameters(), x)