# Copyright 2023 DeepMind Technologies Limited
#
# 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
#
#    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.
# ==============================================================================

"""Manages datasets for the transformer code."""
# From haiku transformer example.
from collections.abc import Iterable, Iterator
import itertools
import random
from typing import NamedTuple, TypeVar

import numpy as np

VOCAB_SIZE = 128  # Number of ASCII code points.
PAD_TOKEN = 0

_T = TypeVar('_T')


class Batch(NamedTuple):
  inputs: np.ndarray  # Integer tokens, shape [B, T].
  targets: np.ndarray  # Integer tokens, shape [B, T].


def repeat(dataset: Iterable[_T]) -> Iterator[_T]:
  return itertools.cycle(dataset)


def shuffle(dataset: Iterator[_T], buffer_size: int) -> Iterator[_T]:
  buffer = [next(dataset) for _ in range(buffer_size)]
  random.shuffle(buffer)
  for item in dataset:
    idx = random.randint(0, buffer_size - 1)  # Inclusive.
    result = buffer[idx]
    buffer[idx] = item
    yield result


def load_ascii_dataset(
    corpus: str,
    *,
    batch_size: int,
    sequence_length: int,
    num_shuffle_batches: int = 10,
) -> Iterator[Batch]:
  """Loads a single-file ASCII dataset in memory."""

  if not corpus.isascii():
    raise ValueError('Loaded corpus is not ASCII.')

  if chr(PAD_TOKEN) in corpus:  # Reserve 0 codepoint for pad token.
    raise ValueError('Corpus must not contain the null byte.')

  # Naively tokenise by taking ASCII codepoints.
  corpus = np.array([ord(c) for c in corpus]).astype(np.int32)
  assert np.max(corpus) < VOCAB_SIZE

  crop_len = sequence_length + 1
  num_batches, remainder = divmod(corpus.size, batch_size * crop_len)
  if remainder:
    corpus = corpus[:-remainder]  # Drop remainder (incomplete) batch.
  ds = corpus.reshape([-1, crop_len])

  if num_batches < num_shuffle_batches:
    raise ValueError(
        f'Only {num_batches} batches in the dataset; consider using a shorter '
        'sequence length or a smaller batch batch size.',
    )

  ds = repeat(ds)
  ds = shuffle(ds, buffer_size=batch_size * num_shuffle_batches)
  while True:
    batch = np.stack([next(ds) for _ in range(batch_size)])
    yield Batch(inputs=batch[:, :-1], targets=batch[:, 1:])

dataset.py
google-deepmind/sam_edge
# Copyright 2023 DeepMind Technologies Limited
#
# 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
#
#    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.
# ==============================================================================

"""Compute the norm of the Hessian of a Jax model."""
import jax
from jax import grad
from jax import jit
from jax import tree_util
import jax.numpy as jnp
import more_tree_utils as mtu

# The size of the offsets used to estimate || H ||_op -- the measurement will be
# bounded above and below by the maximum of minimum of || H ||_op in a ball of
# this size centered at H
DIFFERENTIATION_STEP_SIZE = 0.001
POWER_ITERATION_EPS = 0.001  # relative error


class CurvatureEstimator:
  """Estimate the curvature of a loss function."""

  def __init__(self, loss, rng, iteration_limit=100):
    self.loss = loss
    self.loss_grad = grad(loss)
    self.rng = rng
    self.iteration_limit = iteration_limit

  def hessian_vector_product(self, params, x, y, vector):
    """Compute the product between the Hessian and a vector.

    Args:
      params: a pytree with the parameters of the model
      x: feature tensor
      y: label tensor
      vector: a pytree of the same shape as params

    Returns:
      The product of the Hessian of the loss with the vector.
    """

    @jit
    def hessian_vector_product_wo_self(params, x, y, vector):
    # doesn't take self as an input, so that it can be jitted
      eps = DIFFERENTIATION_STEP_SIZE

      plus_location = tree_util.tree_map(lambda theta, v: theta + eps*v,
                                         params,
                                         vector)
      plus_gradient = self.loss_grad(plus_location, x, y)
      minus_location = tree_util.tree_map(lambda theta, v: theta - eps*v,
                                          params,
                                          vector)
      minus_gradient = self.loss_grad(minus_location, x, y)
      # pylint: disable=g-long-lambda
      return tree_util.tree_map(lambda x, y: ((x - y)/
                                              (2.0*eps
                                               + mtu.NORMALIZING_EPS)),
                                plus_gradient,
                                minus_gradient)

    return hessian_vector_product_wo_self(params, x, y, vector)

  def curvature_and_direction(self, params, x, y):
    self.rng, subkey = jax.random.split(self.rng)
    v = mtu.get_random_direction(subkey, params)

    return self.curvature_and_direction_with_start(params, x, y, v)

  def curvature_and_direction_with_start(self, params, x, y, v):
    """Compute the curvature of the loss with respect to x and y.

    Args:
      params: a pytree with the parameters of the model
      x: feature tensor
      y: label tensor
      v: a pytree of the same shape as params, to use as a warm start

    Returns:
      The principal eigenvalue of the Hessian (which could be negative),
      and the principal eigenvector.
    """

    @jit
    def iteration(old_v):
      # perform hessian-vector product
      product = self.hessian_vector_product(params, x, y, old_v)

      # calculate the norm
      product_norm = mtu.get_vector_norm(product)

      # normalize the product to get a new iterate v
      new_v = jax.tree_util.tree_map(lambda w: w/product_norm, product)

      return new_v, product, product_norm

    last_norm = None
    current_norm = 1.0
    t = 0
    while (((last_norm is None)
            or ((abs(current_norm - last_norm)
                 > current_norm*POWER_ITERATION_EPS)
                and abs(current_norm) > POWER_ITERATION_EPS))
           and (t < self.iteration_limit)):
      last_norm = current_norm
      v, product, current_norm = iteration(v)
      t += 1

    leaf_dots_tree = tree_util.tree_map(lambda x, y: jnp.sum(x*y), v, product)
    leaf_dots = tree_util.tree_leaves(leaf_dots_tree)
    return jnp.sum(jnp.array(leaf_dots)), v

  def hessian_operator_norm(self, params, x, y):
    mu, _ = self.curvature_and_direction(params, x, y)
    return abs(mu)

  # * not vectorized (at least among the components), for now
  # * to save memory, only return the principal eigenvector
  def hessian_top_eigenvalues(self, params, x, y, k):
    """Compute the k top eigenvalues of the Hessian.

    Args:
      params: A pytree with the parameters of the model
      x: feature tensor
      y: label tensor
      k: number of eigenvalues to return

    Returns:
      A jnp array with the eigenvalues, and pytree of the same shape
      as params with the principal eigenvector.

    Note: assumes that tghe top eigenvalues are positive.
    """

    component_list = list()
    for _ in range(k):
      self.rng, subkey = jax.random.split(self.rng)
      component_list.append(mtu.get_random_direction(subkey, params))
    component_list = mtu.get_orthonormal_basis(component_list)

    @jit
    def iteration(old_components):
      products = list()
      product_norms = list()
      for i in range(k):
        # perform hessian-vector product
        products.append(self.hessian_vector_product(params, x, y,
                                                    old_components[i]))
        product_norms.append(mtu.get_vector_norm(products[i]))

      return mtu.get_orthonormal_basis(products), jnp.array(product_norms)

    last_norms = None
    current_norms = jnp.ones(k)
    t = 0
    while (((last_norms is None)
            or ((jnp.linalg.norm(current_norms - last_norms)
                 > jnp.linalg.norm(current_norms)*POWER_ITERATION_EPS)
                and jnp.linalg.norm(current_norms) > POWER_ITERATION_EPS))
           and (t < self.iteration_limit)):
      last_norms = current_norms
      component_list, current_norms = iteration(component_list)
      t += 1

    # sort in decreasing order
    indices_by_norms = jnp.argsort(-current_norms)

    return current_norms[indices_by_norms], component_list[indices_by_norms[0]]

hessian_norm.py
google-deepmind/sam_edge
# Copyright 2023 DeepMind Technologies Limited
#
# 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
#
#    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.
# ==============================================================================

"""Perform image classification experiments regarding SAM and the edge of stability."""
# Modification from haiku/examples/transformer/train.py
import argparse
import time
from typing import Any, Sequence

from flax import linen as nn
from jax import jit
import jax.numpy as jnp
import jax.random as jrandom
import more_tree_utils as mtu
import optax
import sam_edge
import tensorflow as tf
import tensorflow_datasets as tfds

parser = argparse.ArgumentParser()

parser.add_argument("--time_limit_in_hours",
                    type=float,
                    default=4.0)

parser.add_argument("--hessian_check_gap",
                    type=float,
                    default=0.25,
                    help=("If not None, the number of seconds "
                          + "between hessian-norm measurements"))

parser.add_argument("--step_size",
                    type=float,
                    default=0.01)

parser.add_argument("--batch_size",
                    type=int,
                    default=60000)

parser.add_argument("--nn_architecture",
                    type=str,
                    default="MLP",
                    choices=["CNN", "MLP"])

parser.add_argument("--dataset",
                    type=str,
                    default="mnist",
                    choices=["mnist", "cifar10"])

parser.add_argument("--mlp_depth",
                    type=int,
                    default=4)

parser.add_argument("--mlp_width",
                    type=int,
                    default=1000)

parser.add_argument("--cnn_num_blocks",
                    type=int,
                    default=2)

parser.add_argument("--cnn_layers_per_block",
                    type=int,
                    default=1)

parser.add_argument("--cnn_feature_multiplier",
                    type=int,
                    default=16)

parser.add_argument("--mini_training_set_num_batches",
                    type=int,
                    default=None,
                    help=("If this is not None, it is the number "
                          + "of batches in a reduced training set"))

parser.add_argument("--mini_test_set_num_batches",
                    type=int,
                    default=None,
                    help=("If this is not None, it is the number "
                          +"of batches in a reduced test set"))

parser.add_argument("--rho",
                    type=float,
                    default=0.0,
                    help=("The parameter rho for SAM -- "
                          + "if rho is 0, SAM is not used"))

parser.add_argument("--eigs_curve_output",
                    type=str,
                    default="/tmp/eigs.pdf")

parser.add_argument("--eigs_se_only_output",
                    type=str,
                    default="/tmp/eigs_se_only.pdf",
                    help=("Output for plotting the eigenvalues "
                          + "of the hessian and the SAM-edge only"))

parser.add_argument("--alignment_curve_output",
                    type=str,
                    default="/tmp/a.pdf")

parser.add_argument("--loss_curve_output",
                    type=str,
                    default="/tmp/ell.pdf")

parser.add_argument("--raw_data_output",
                    type=str,
                    default="/tmp/raw.txt")

parser.add_argument("--num_principal_components",
                    type=int,
                    default=1)

args = parser.parse_args()

SEED_FACTOR = 100000
SHUFFLE_BATCH_SIZE = 1024

ModuleDef = Any


class CNN(nn.Module):
  """A standard convolutional neural network."""
  conv_defs: Sequence[ModuleDef]
  dense_def: ModuleDef

  @nn.compact
  def __call__(self, x):
    h = x
    for block_conv_defs in self.conv_defs:
      for conv_def in block_conv_defs:
        h = conv_def(h)
        h = nn.relu(h)
        h = nn.LayerNorm()(h)
      h = nn.max_pool(h, window_shape=(2, 2), strides=(2, 2))
    h = h.reshape((h.shape[0], -1))  # flatten
    h = self.dense_def(h)
    return h


class MLP(nn.Module):
  """A standard multi-layer perceptron."""
  input_to_hidden_def: ModuleDef
  hidden_to_hidden_def: ModuleDef
  output_def: ModuleDef

  @nn.compact
  def __call__(self, x):
    h = x.reshape((x.shape[0], -1))  # flatten
    h = self.input_to_hidden_def(h)
    h = nn.relu(h)
    for _ in range(args.mlp_depth-2):
      h = self.hidden_to_hidden_def(h)
      h = nn.relu(h)
    h = self.output_def(h)
    return h


def test_error_fn(params_, model_, batches):
  """Compute test error.

  Args:
    params_: parameters of the model
    model_: the model
    batches: test data

  Returns:
    Test error.
  """

  @jit
  # targets are not one-hot encoded in the test data
  def batch_error(images, targets):
    predicted_class = jnp.argmax(model_.apply(params_, images), axis=1)
    return jnp.mean(predicted_class != targets)

  sum_batch_errors = 0.0
  num_batches = 0
  for x, y in batches:
    sum_batch_errors += batch_error(x, y)
    num_batches += 1
  return sum_batch_errors/num_batches

# ================================
# Start of main routine
print("running")

# Ensure TF does not see GPU and grab all GPU memory.
tf.config.set_visible_devices([], device_type="GPU")

data_dir = "/tmp/" + args.dataset
if args.dataset == "mnist":
  num_training_examples = 60000
  num_test_examples = 10000
  image_size = 28
  num_channels = 1
  num_classes = 10
elif args.dataset == "cifar10":
  num_training_examples = 50000
  num_test_examples = 10000
  image_size = 32
  num_channels = 3
  num_classes = 10

rng = jrandom.PRNGKey(int(SEED_FACTOR*time.time()) % SEED_FACTOR)
gn = nn.initializers.glorot_normal()
if args.nn_architecture == "CNN":
  # pylint: disable=g-complex-comprehension
  model = CNN(conv_defs=[[nn.Conv(features=(2**i*args.cnn_feature_multiplier),
                                  kernel_size=(3, 3),
                                  kernel_init=gn,
                                  padding="SAME")
                          for j in range(args.cnn_layers_per_block)]
                         for i in range(args.cnn_num_blocks)],
              dense_def=nn.Dense(features=num_classes,
                                 kernel_init=gn)
              )
  num_linear_layers = args.cnn_layers_per_block * args.cnn_num_blocks + 1
else:
  model = MLP(input_to_hidden_def=nn.Dense(features=args.mlp_width,
                                           kernel_init=gn),
              hidden_to_hidden_def=nn.Dense(features=args.mlp_width,
                                            kernel_init=gn),
              output_def=nn.Dense(features=num_classes,
                                  kernel_init=gn))
  num_linear_layers = args.mlp_depth
rng, subkey = jrandom.split(rng)
params = model.init(subkey,
                    jnp.ones([args.batch_size,
                              image_size,
                              image_size,
                              num_channels]))
print("parameter count: {}".format(mtu.count_parameters(params)))


def get_train_batches():
  """Get training data."""
  # as_supervised=True gives us the (image, label) as a tuple instead of a dict
  ds = tfds.load(name=args.dataset,
                 split="train",
                 as_supervised=True,
                 data_dir=data_dir)
  if args.mini_training_set_num_batches:
    ds = ds.take(args.mini_training_set_num_batches*args.batch_size)
  # pylint: disable=g-long-lambda
  ds = ds.map(lambda x, y:
              (tf.cast(x, dtype=tf.float32)/256.0, tf.one_hot(y, num_classes)))
  ds = ds.batch(args.batch_size, drop_remainder=True)
  ds = ds.prefetch(tf.data.AUTOTUNE).repeat()
  return tfds.as_numpy(ds)


def get_test_batches(num_available_test_examples):
  """Get test data."""
  ds = tfds.load(name=args.dataset,
                 split="test",
                 as_supervised=True,
                 data_dir=data_dir)
  if args.mini_test_set_num_batches:
    reduced_num_test_examples = (args.mini_test_set_num_batches
                                 *args.batch_size)
    ds = ds.take(num_test_examples)
  else:
    reduced_num_test_examples = num_available_test_examples
  ds = ds.map(lambda x, y: (tf.cast(x, dtype=tf.float32)/256.0, y))
  batch_size = min(reduced_num_test_examples, args.batch_size)
  ds = ds.batch(batch_size, drop_remainder=True)
  ds = ds.prefetch(tf.data.AUTOTUNE)
  return tfds.as_numpy(ds)


# pylint: disable=unused-argument
def loss(x, logits, y):
  return optax.squared_error(logits, y)

train_batches = get_train_batches()
test_batches = get_test_batches(num_test_examples)

# center the data
u, _ = next(iter(train_batches))
mu = jnp.mean(u, axis=0, keepdims=True)
train_batches = ((x - mu, y) for (x, y) in train_batches)
test_batches = ((x - mu, y) for (x, y) in test_batches)


params = sam_edge.train(params,
                        model,
                        loss,
                        train_batches,
                        args.step_size,
                        args.rho,
                        args.hessian_check_gap,
                        args.eigs_curve_output,
                        args.eigs_se_only_output,
                        args.alignment_curve_output,
                        args.loss_curve_output,
                        args.raw_data_output,
                        args.num_principal_components,
                        args.time_limit_in_hours,
                        rng)

test_err = test_error_fn(params, model, test_batches)
print("==============")
print("Test error: {}".format(test_err))


image_classification.py
google-deepmind/sam_edge
# Copyright 2023 DeepMind Technologies Limited
#
# 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
#
#    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.
# ==============================================================================

"""Misc. utilities for pytrees."""
import jax
from jax import jit
from jax import tree_util
import jax.numpy as jnp

NORMALIZING_EPS = 1e-5


@jit
def normalize(t):
  norm = get_vector_norm(t)
  return tree_util.tree_map(lambda t_leaf: t_leaf/(norm + NORMALIZING_EPS), t)


@jit
def project_out_and_normalize(s, t):
  """Project out and normalize.

  Args:
    s: pytree
    t: another pytree

  Returns:
    The pytree obtained by projecting the flattening of
    t onto the flattening of s, normalizing the result,
    and then reshaping into a pytree.
  """

  s_dot_t = get_tree_dot(s, t)
  # pylint: disable=g-long-lambda
  new_part = tree_util.tree_map(lambda s_leaf, t_leaf:
                                t_leaf - s_dot_t * s_leaf, s, t)
  return normalize(new_part)


def get_orthonormal_basis(t_list):
  k = len(t_list)
  t_list[0] = normalize(t_list[0])
  for i in range(k):
    for j in range(i+1, k):
      t_list[j] = project_out_and_normalize(t_list[i], t_list[j])
  return t_list


@jit
def get_vector_norm(t):
  squared_norms = tree_util.tree_map(lambda x: jnp.sum(x*x), t)
  return jnp.sqrt(jnp.sum(jnp.array(tree_util.tree_leaves(squared_norms))))


def count_parameters(t):
  leaf_parameter_counts = tree_util.tree_map(lambda x: x.size, t)
  return jnp.sum(jnp.array(tree_util.tree_leaves(leaf_parameter_counts)))


@jit
def get_tree_dot(s, t):
  leaf_dots = tree_util.tree_map(lambda si, ti: jnp.sum(si*ti), s, t)
  return jnp.sum(jnp.array(tree_util.tree_leaves(leaf_dots)))


@jit
def get_alignment(s, t):
  return jnp.abs(get_tree_dot(s, t))/(get_vector_norm(s)*get_vector_norm(t))


@jit
def get_random_direction(rng_key, t):
  """Sample a unit length pytree.

  Args:
    rng_key: RNG key
    t: a pytree

  Returns:
    A pytree with the same shape as t, whose
    leaves are collectively sampled uniformly
    at random from the unit ball.
  """
  def sample_at_leaf(sub_key, shape):
    return jax.random.normal(sub_key, shape)

  flat_t, treedef = tree_util.tree_flatten(t)
  leaf_shapes = tree_util.tree_map(lambda x: x.shape, flat_t)
  rng_keys = jax.random.split(rng_key, len(leaf_shapes))

  new_leaves = [sample_at_leaf(rng_keys[i], leaf_shapes[i])
                for i in range(len(leaf_shapes))]
  new_leaves = normalize(new_leaves)
  return tree_util.tree_unflatten(treedef, new_leaves)

more_tree_utils.py
google-deepmind/sam_edge
# Copyright 2023 DeepMind Technologies Limited
#
# 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
#
#    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.
# ==============================================================================

"""Perform experiments regarding SAM and the edge of stability."""
import dataclasses
import math
import time

import hessian_norm
from jax import grad
from jax import jit
from jax import tree_util
import jax.numpy as jnp
import matplotlib.pyplot as plt
import more_tree_utils as mtu

EPSILON = 1e-5
DPI = 300

# pylint: disable=anomalous-backslash-in-string


def train(params,
          model,
          loss,
          train_batches,
          step_size,
          rho,  # for SAM -- if rho is 0.0, SAM is not used
          hessian_check_gap,
          eigs_curve_filename,
          eigs_se_only_filename,
          alignment_curve_filename,
          loss_curve_filename,
          raw_data_filename,
          num_principal_comps,
          time_limit_in_hours,
          rng):
  """Train a model using SAM, and plot statistics.

  Args:
    params: parameters of the model
    model: the model
    loss: the loss function
    train_batches: training data
    step_size: learning rate
    rho: distance uphill to evaluate the gradient
    hessian_check_gap: time in seconds hessian evaluation
    eigs_curve_filename: name of PDF file for eigenvalue/edge plots
    eigs_se_only_filename: name of PDF file for eigenvalue/edge plots
            without 2/eta
    alignment_curve_filename: name of PDF file for plots of alignments
    loss_curve_filename: name of PDF file for loss curves
    raw_data_filename: name of file for raw data
    num_principal_comps: number of the principal components of the
            hessian to evaluate
    time_limit_in_hours: Time limit
    rng: key

  Returns:
    final parameters
  """

  @jit
  def loss_by_params(params, x_batched, y_batched):
    preds = model.apply(params, x_batched)
    return jnp.mean(loss(x_batched, preds, y_batched))

  @jit
  def sam_neighbor(params, x, y):
    grads = grad(loss_by_params)(params, x, y)
    norm = mtu.get_vector_norm(grads)
    return tree_util.tree_map(lambda p, g: p + rho * g/(norm + EPSILON),
                              params,
                              grads)

  @jit
  def update(params, x, y, eta):
    if rho > 0.0:
      grad_location = sam_neighbor(params, x, y)
    else:
      grad_location = params
    grads = grad(loss_by_params)(grad_location, x, y)
    return tree_util.tree_map(lambda p, g: p - eta * g,
                              params,
                              grads)

  @jit
  def get_sam_gradient(params, x, y):
    grad_location = sam_neighbor(params, x, y)
    return grad(loss_by_params)(grad_location, x, y)

  eta = step_size

  @dataclasses.dataclass
  class PlotData:
    # pylint: disable=g-bare-generic
    training_times: list
    eigenvalues: list
    sam_edges: list
    g_alignments: list
    sg_alignments: list
    training_losses: list
  plot_data = PlotData(list(),
                       list(),
                       list(),
                       list(),
                       list(),
                       list())
  if num_principal_comps > 1:
    for i in range(num_principal_comps):
      plot_data.eigenvalues.append(list())

  ce = hessian_norm.CurvatureEstimator(loss_by_params, rng)

  print("starting training", flush=True)
  start_time = time.time()
  last_hessian_check = start_time
  time_limit = 3600*time_limit_in_hours
  this_loss = None
  for x, y in train_batches:
    if ((time.time() > start_time + time_limit)
        or (this_loss and jnp.isnan(this_loss))):
      break

    if (hessian_check_gap and
        (time.time() > last_hessian_check + 3600.0*hessian_check_gap)):
      original_gradient = grad(loss_by_params)(params, x, y)
      sam_gradient = get_sam_gradient(params, x, y)
      if num_principal_comps == 1:
        curvature, principal_dir = ce.curvature_and_direction(params, x, y)
        this_hessian_norm = jnp.abs(curvature)
      else:
        print("calculating principal components", flush=True)
        eigs, principal_dir = ce.hessian_top_eigenvalues(params, x, y,
                                                         num_principal_comps)
        print("done calculating principal components", flush=True)
        this_hessian_norm = eigs[0]
      grad_hessian_alignment = mtu.get_alignment(original_gradient,
                                                 principal_dir)
      samgrad_hessian_alignment = mtu.get_alignment(sam_gradient,
                                                    principal_dir)
      training_time = time.time() - start_time
      original_gradient_norm = mtu.get_vector_norm(original_gradient)
      this_loss = loss_by_params(params, x, y)

      if rho == 0.0:
        sam_edge = 2.0/eta
      else:
        sam_edge = ((original_gradient_norm/(2.0*rho))
                    *(math.sqrt(1.0
                                + 8.0*rho/(eta*original_gradient_norm))
                      - 1.0))
      print("--------------", flush=True)
      formatting_string = ("Norm: {}, "
                           + "2/eta: {}, "
                           + "sam_edge: {}, "
                           + "|| g || = {}, "
                           + "loss = {}, "
                           + "g_alignment = {}, "
                           + "sg_alignment = {}")
      print(formatting_string.format(this_hessian_norm,
                                     2.0/eta,
                                     sam_edge,
                                     original_gradient_norm,
                                     this_loss,
                                     grad_hessian_alignment,
                                     samgrad_hessian_alignment, flush=True))
      if num_principal_comps > 1:
        print("eigs = {}".format(eigs, flush=True))
      if eigs_curve_filename or eigs_se_only_filename:
        plot_data.training_times.append(training_time)
        if num_principal_comps == 1:
          plot_data.eigenvalues.append(this_hessian_norm)
        else:
          for i in range(num_principal_comps):
            plot_data.eigenvalues[i].append(eigs[i])
        plot_data.sam_edges.append(sam_edge)
        plot_data.g_alignments.append(grad_hessian_alignment)
        plot_data.sg_alignments.append(samgrad_hessian_alignment)
        plot_data.training_losses.append(this_loss)
      if raw_data_filename:
        with open(raw_data_filename, "a") as raw_data_file:
          columns = [training_time,
                     this_hessian_norm,
                     2.0/eta,
                     sam_edge,
                     original_gradient_norm,
                     this_loss,
                     grad_hessian_alignment,
                     samgrad_hessian_alignment]
          format_string = "{} "*(len(columns)-1) + "{}\n"
          raw_data_file.write(format_string.format(*columns))
      last_hessian_check = time.time()

    params = update(params, x, y, eta)

  if (plot_data.sam_edges
      and (not jnp.isnan(jnp.array(plot_data.training_losses)).any())):
    if eigs_curve_filename:
      plt.figure()
      if num_principal_comps == 1:
        max_y = max(2.0/eta,
                    max(plot_data.eigenvalues),
                    max(plot_data.sam_edges))
      else:
        max_y = max(2.0/eta, max(plot_data.sam_edges))
        for i in range(num_principal_comps):
          max_y = max(max_y, max(plot_data.eigenvalues[i]))
      plt.ylim(0, 1.1*max_y)
      if num_principal_comps == 1:
        plt.plot(plot_data.training_times,
                 plot_data.eigenvalues,
                 color="b",
                 label="$|| H ||_{op}$")
      else:
        for i in range(num_principal_comps):
          plt.plot(plot_data.training_times,
                   plot_data.eigenvalues[i],
                   color="b",
                   label="$\lambda_{}$".format(i+1))
      plt.plot(plot_data.training_times,
               plot_data.sam_edges,
               color="g",
               label="SAM edge")
      plt.axhline(2.0/eta, color="m", label="$2/\eta$")
      plt.legend()
      plt.savefig(eigs_curve_filename, format="pdf", dpi=DPI)
    if eigs_se_only_filename:
      plt.figure()
      if num_principal_comps == 1:
        max_y = max(max(plot_data.eigenvalues), max(plot_data.sam_edges))
      else:
        max_y = max(plot_data.sam_edges)
        for i in range(num_principal_comps):
          max_y = max(max_y, max(plot_data.eigenvalues[i]))
      plt.ylim(0, 1.1*max_y)
      if num_principal_comps == 1:
        plt.plot(plot_data.training_times,
                 plot_data.eigenvalues,
                 color="b",
                 label="$|| H ||_{op}$")
      else:
        for i in range(num_principal_comps):
          plt.plot(plot_data.training_times,
                   plot_data.eigenvalues[i],
                   color="b",
                   label="$\lambda_{}$".format(i+1))
      plt.plot(plot_data.training_times,
               plot_data.sam_edges,
               color="g",
               label="SAM edge")
      plt.legend()

      plt.savefig(eigs_se_only_filename, format="pdf", dpi=DPI)
    if alignment_curve_filename:
      plt.figure()
      plt.ylim(0, 1.1*max(jnp.max(jnp.array(plot_data.g_alignments)),
                          jnp.max(jnp.array(plot_data.sg_alignments))))
      plt.plot(plot_data.training_times,
               plot_data.g_alignments,
               color="b",
               label="original gradient alignments")
      plt.plot(plot_data.training_times,
               plot_data.sg_alignments,
               color="g",
               label="SAM gradient alignments")
      plt.legend()

      plt.savefig(alignment_curve_filename, format="pdf", dpi=DPI)
    if loss_curve_filename:
      plt.figure()
      plt.ylim(0, 1.1*jnp.max(jnp.array(plot_data.training_losses)))
      plt.plot(plot_data.training_times,
               plot_data.training_losses,
               color="b",
               label="training loss")
      plt.legend()

      plt.savefig(loss_curve_filename, format="pdf", dpi=DPI)

  return params

sam_edge.py
google-deepmind/sam_edge
# Copyright 2023 DeepMind Technologies Limited
#
# 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
#
#    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.
# ==============================================================================

"""Run experiments on SAM and the edge of stability using a transformer model."""
# This is a modification of haiku/examples/transformer/train.py
import argparse
import time
from typing import Any, MutableMapping

import dataset
import haiku as hk
import jax
from jax import jit
import jax.numpy as jnp
import jax.random as jrandom
import more_tree_utils as mtu
import sam_edge
import transformer_model

# pylint: disable=bad-continuation

parser = argparse.ArgumentParser()

parser.add_argument("--training_data",
                    type=str,
                    required=True)

parser.add_argument("--test_data",
                    type=str,
                    required=True)

parser.add_argument("--batch_size",
                    type=int,
                    default=128)

parser.add_argument("--time_limit_in_hours",
                    type=float,
                    default=1.0)

parser.add_argument("--hessian_check_gap",
                    type=float,
                    help=("If not None, the number of hours "
                          + "between checks of the hessian norm"))

parser.add_argument("--num_layers",
                    type=int,
                    default=6)

parser.add_argument("--num_heads",
                    type=int,
                    default=8)

parser.add_argument("--key_size",
                    type=int,
                    default=32)

parser.add_argument("--model_size",
                    type=int,
                    default=128)

parser.add_argument("--sequence_length",
                    type=int,
                    default=64)

parser.add_argument("--step_size",
                    type=float,
                    default=0.1)

parser.add_argument("--rho",
                    type=float,
                    default=0.0,
                    help="SAM offset parameter -- if rho is 0, SAM is not used")

parser.add_argument("--eigs_curve_output",
                    type=str,
                    default="/tmp/eigs.pdf")

parser.add_argument("--eigs_se_only_output",
                    type=str,
                    default="/tmp/eigs_se_only.pdf",
                    help=("Output for plotting the eigenvalues of "
                          + "the hessian and the SAM-edge only"))

parser.add_argument("--alignment_curve_output",
                    type=str,
                    default="/tmp/a.pdf")

parser.add_argument("--loss_curve_output",
                    type=str,
                    default="/tmp/ell.pdf")

parser.add_argument("--raw_data_output",
                    type=str,
                    default="/tmp/raw.txt")

parser.add_argument("--num_principal_components",
                    type=int,
                    default=1)

args = parser.parse_args()

SEED_FACTOR = 100000
VOCAB_SIZE = 128
PAD_TOKEN = 0

_Batch = dataset.Batch
_Metrics = MutableMapping[str, Any]


def loss_fn(x, logits, y):
  """Computes the (scalar) language modelling loss on `data` w.r.t. params."""
  log_probs = jax.nn.log_softmax(logits)  # [B, T, V]
  onehot_y = jax.nn.one_hot(y, VOCAB_SIZE)
  log_likelihood = jnp.sum(onehot_y * log_probs, axis=-1)  # [B, T]

  # Loss is the average negative log-likelihood per (non-masked) token.
  mask = jnp.not_equal(x, dataset.PAD_TOKEN)  # [B, T]
  return -jnp.sum(log_likelihood * mask) / jnp.sum(mask)  # []

# Create the datasets
print("Generating training data")
with open(args.training_data, mode="r") as file:
  train_dataset = dataset.load_ascii_dataset(
    corpus=file.read(),
    batch_size=args.batch_size,
    sequence_length=args.sequence_length
  )
print("Generating test data")
with open(args.test_data, mode="r") as file:
  raw_test_data = file.read()
  num_test_batches = (len(raw_test_data)//
                      (args.batch_size*args.sequence_length))
  test_dataset = dataset.load_ascii_dataset(
    corpus=raw_test_data,
    batch_size=args.batch_size,
    sequence_length=args.sequence_length
  )


def test_loss_fn(params_, lm, test_dataset_, num_test_batches_):
  """Computes the test loss."""
  @jit
  def batch_error(x, y):
    logits = lm.apply(params_, x)
    return loss_fn(x, logits, y)

  total_loss = 0.0
  num_batches = 0
  for data_ in test_dataset_:
    total_loss += batch_error(data_.inputs, data_.targets)
    num_batches += 1
    if num_batches == num_test_batches_:
      break
  return total_loss/num_batches

print("Initializing model")


def forward_pass(x):
  lm = transformer_model.LanguageModel(
    model_size=args.model_size,
    vocab_size=VOCAB_SIZE,
    pad_token=PAD_TOKEN,
    transformer=transformer_model.Transformer(
      num_heads=args.num_heads,
      num_layers=args.num_layers,
      attn_size=args.key_size,
    ))
  return lm(x)

forward_pass = hk.transform(forward_pass)
forward_pass = hk.without_apply_rng(forward_pass)
rng = jax.random.PRNGKey(int(SEED_FACTOR*time.time()) % SEED_FACTOR)
data = next(train_dataset)
rng, subkey = jrandom.split(rng)
params = forward_pass.init(subkey, data.inputs)
parameter_count = mtu.count_parameters(params)
print(f"  {parameter_count} parameters")

print("Training")
params = sam_edge.train(params,
                        forward_pass,
                        loss_fn,
                        train_dataset,
                        args.step_size,
                        args.rho,
                        args.hessian_check_gap,
                        args.eigs_curve_output,
                        args.eigs_se_only_output,
                        args.alignment_curve_output,
                        args.loss_curve_output,
                        args.raw_data_output,
                        args.num_principal_components,
                        args.time_limit_in_hours,
                        rng)

test_loss = test_loss_fn(params,
                         forward_pass,
                         test_dataset,
                         num_test_batches)
print("=======================")
print(f"test_loss: {test_loss}")

transformer.py
google-deepmind/sam_edge
# Copyright 2023 DeepMind Technologies Limited
#
# 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
#
#    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.
# ==============================================================================

"""Library with Transformer model definition."""
# Modification of haiku/examples/transformer/model.py
import dataclasses
from typing import Optional

import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np


def _layer_norm(x: jax.Array) -> jax.Array:
  """Applies a unique LayerNorm to `x` with default settings."""
  ln = hk.LayerNorm(axis=-1, create_scale=True, create_offset=True)
  return ln(x)


@dataclasses.dataclass
class Transformer(hk.Module):
  """A transformer stack."""

  num_heads: int  # Number of attention heads.
  num_layers: int  # Number of transformer (attention + MLP) layers to stack.
  attn_size: int  # Size of the attention (key, query, value) vectors.
  widening_factor: int = 4  # Factor by which the MLP hidden layer widens.
  name: Optional[str] = None  # Optional identifier for the module.

  def __call__(
      self,
      embeddings: jax.Array,  # [B, T, D]
      mask: jax.Array,  # [B, T]
  ) -> jax.Array:  # [B, T, D]
    """Transforms input embedding sequences to output embedding sequences."""

    initializer = hk.initializers.VarianceScaling(2 / self.num_layers)
    _, seq_len, model_size = embeddings.shape

    # Compute causal mask for autoregressive sequence modelling.
    mask = mask[:, None, None, :]  # [B, H=1, T'=1, T]
    causal_mask = np.tril(np.ones((1, 1, seq_len, seq_len)))  # [B=1, H=1, T, T]
    mask = mask * causal_mask  # [B, H=1, T, T]

    h = embeddings
    for _ in range(self.num_layers):
      # First the attention block.
      attn_block = hk.MultiHeadAttention(
          num_heads=self.num_heads,
          key_size=self.attn_size,
          model_size=model_size,
          w_init=initializer,
      )
      h_norm = _layer_norm(h)
      h_attn = attn_block(h_norm, h_norm, h_norm, mask=mask)
      h = h + h_attn

      # Then the dense block.
      dense_block = hk.Sequential([
          hk.Linear(self.widening_factor * model_size, w_init=initializer),
          jax.nn.gelu,
          hk.Linear(model_size, w_init=initializer),
      ])
      h_norm = _layer_norm(h)
      h_dense = dense_block(h_norm)
      h = h + h_dense

    return _layer_norm(h)


@dataclasses.dataclass
class LanguageModel(hk.Module):
  """An autoregressive transformer-based language model."""

  transformer: Transformer
  model_size: int  # Embedding size.
  vocab_size: int  # Size of the vocabulary.
  pad_token: int  # Identity of the padding token (used for masking inputs).
  name: Optional[str] = None  # Optional identifier for the module.

  def __call__(
      self,
      tokens: jax.Array,  # Batch of sequences of input tokens, shape [B, T].
  ) -> jax.Array:  # Batch of sequences of output token logits, shape [B, T, V].
    """Forward pass, producing a sequence of logits."""
    input_mask = jnp.greater(tokens, self.pad_token)
    unused_batch_size, seq_len = tokens.shape

    # Embed the input tokens and positions.
    embed_init = hk.initializers.TruncatedNormal(stddev=0.02)
    token_embedding_map = hk.Embed(
        self.vocab_size, embed_dim=self.model_size, w_init=embed_init)
    token_embeddings = token_embedding_map(tokens)
    positional_embeddings = hk.get_parameter(
        'positional_embeddings', [seq_len, self.model_size], init=embed_init)
    input_embeddings = token_embeddings + positional_embeddings  # [B, T, D]

    # Run the transformer over the inputs.
    embeddings = self.transformer(input_embeddings, input_mask)  # [B, T, D]

    # Decode the embeddings (here, we use untied weights).
    return hk.Linear(self.vocab_size)(embeddings)  # [B, T, V]

transformer_model.py
google-deepmind/sam_edge