Add different nonlinearities on std head.

PiperOrigin-RevId: 691152711

Author: xingyousong

optformer/embed_then_regress/configs.py

@@ -20,6 +20,7 @@
 from typing import Callable
 from flax import linen as nn
 import jax
+import jaxtyping as jt
 import optax
 from optformer.embed_then_regress import icl_transformer
 from optformer.t5x import embedders
@@ -76,18 +77,45 @@ class ModelConfig:
   nhead: int = 16
   dropout: float = 0.1
   num_layers: int = 8
+  std_transform: str = 'exp'
 
   def create_model(
       self, embedder_config: T5EmbedderConfig
   ) -> icl_transformer.ICLTransformer:
 
     kwargs = dataclasses.asdict(self)
     embedder_factory = embedder_config.create_embedder_factory()
+    std_transform_fn = self.create_std_transform_fn()
 
     return icl_transformer.ICLTransformer(
-        embedder_factory=embedder_factory, **kwargs
+        std_transform_fn=std_transform_fn,
+        embedder_factory=embedder_factory,
+        **kwargs,
     )
 
+  def create_std_transform_fn(
+      self,
+  ) -> Callable[[jt.Float[jax.Array, '*A']], jt.Float[jax.Array, '*A']]:
+    """Creates std transform function."""
+    if self.std_transform == 'exp':
+      return jax.numpy.exp
+    elif self.std_transform == 'exp10':
+      return lambda x: jax.numpy.exp(10.0 * x)
+    elif self.std_transform == 'softplus':
+      return jax.nn.softplus
+    elif self.std_transform == 'softplus10':
+      return lambda x: jax.nn.softplus(10.0 * x)
+    elif self.std_transform == 'abs':
+      return jax.numpy.abs
+    elif self.std_transform == 'abs10':
+      return lambda x: jax.numpy.abs(10.0 * x)
+    elif self.std_transform == 'shifted_relu':
+      return lambda x: jax.nn.relu(x + 1.0)
+    elif self.std_transform == 'shifted_relu10':
+      return lambda x: jax.nn.relu(10.0 * x + 1.0)
+    else:
+      raise ValueError(f'Unknown std_transform: {self.std_transform}')
+
 
 @dataclasses.dataclass
 class TrainingConfig:

optformer/embed_then_regress/icl_transformer.py

@@ -30,6 +30,8 @@
 # order to use the same learning rate.
 default_kernel_init = nn.initializers.truncated_normal(stddev=0.02)
 Dense = functools.partial(nn.Dense, kernel_init=default_kernel_init)
+EPS = 1e-7
+AnyTensor = jt.Float[jax.Array, '*A']
 
 
 class Block(nn.Module):
@@ -92,7 +94,7 @@ class ICLTransformer(nn.Module):
   nhead: int  # H
   dropout: float
   num_layers: int
-
+  std_transform_fn: Callable[[AnyTensor], AnyTensor]
   embedder_factory: Callable[[], nn.Module]  # __call__: [B, T] -> [B, D]
 
   def setup(self):
@@ -166,8 +168,8 @@ def __call__(
     for layer in self.encoder_layers:
       out = layer(out, mask, deterministic, rng)
 
-    mean, log_std = jnp.split(self.mean_logstd_head(out), 2, axis=-1)  # [B L 1]
-    std = jnp.exp(log_std)
+    mean, std = jnp.split(self.mean_logstd_head(out), 2, axis=-1)  # [B L 1]
+    std = self.std_transform_fn(self.std_transform)(std) + EPS
 
     mean = jnp.squeeze(mean, axis=-1)
     std = jnp.squeeze(std, axis=-1)

optformer/embed_then_regress/train.py

@@ -31,7 +31,6 @@
 
 
 Scalar = jnp.ndarray | np.ndarray | float
-EPS = 1e-7
 
 
 def multi_gpu() -> bool:
@@ -98,7 +97,7 @@ def loss_fn(
   """Loss function with metrics."""
   # pylint: disable=invalid-name
   mean, std = model.apply(params, deterministic=not training, rng=rng, **batch)
-  nlogprob = -jax.scipy.stats.norm.logpdf(batch['y'], mean, std + EPS)  # [B, L]
+  nlogprob = -jax.scipy.stats.norm.logpdf(batch['y'], mean, std)  # [B, L]
 
   # Only compute loss over target ys. Mask is BxL where True denotes context
   # token and False otherwise.