Changes

Author: cg123

mergekit/scripts/tokensurgeon.py

@@ -88,6 +88,8 @@ class TokenSurgeonOptions(BaseModel):
     cosine_similarity: bool = False
     subword_method: SubwordMethod = SubwordMethod.MEAN
     batch_size: Optional[int] = None
+    new_vocab_noise: Optional[float] = None
+    new_vocab_scale: Optional[float] = None
 
 
 def get_arch_info(
@@ -489,6 +491,10 @@ def build_embedding_matrix(
                 token_basis=token_basis,
                 options=options,
             )
+            if options.new_vocab_noise:
+                new_embeds += torch.randn_like(new_embeds) * options.new_vocab_noise
+            if options.new_vocab_scale:
+                new_embeds *= options.new_vocab_scale
             for ne_idx, token in enumerate(
                 new_tokens[base_idx : base_idx + batch_size]
             ):
@@ -592,6 +598,22 @@ class AllowMatch(enum.Enum):
     help="Filter out poorly trained tokens",
     show_default=True,
 )
+@click.option(
+    "--new-vocab-noise",
+    "-nvn",
+    type=float,
+    default=None,
+    help="Add gaussian noise to new vocab embeddings",
+    show_default=True,
+)
+@click.option(
+    "--new-vocab-scale",
+    "-nvs",
+    type=float,
+    default=None,
+    help="Scale computed new vocab embeddings by this factor",
+    show_default=True,
+)
 @add_merge_options
 def main(
     model: str,
@@ -607,6 +629,8 @@ def main(
     prefix_match: str,
     byte_match: str,
     magikarp: bool,
+    new_vocab_noise: Optional[float],
+    new_vocab_scale: Optional[float],
     merge_options: MergeOptions,
 ):
     merge_options.apply_global_options()
@@ -622,6 +646,8 @@ def main(
         weight_scheme=WeightingScheme(weight_scheme),
         subword_method=SubwordMethod(subword_method),
         batch_size=batch_size,
+        new_vocab_noise=new_vocab_noise,
+        new_vocab_scale=new_vocab_scale,
     )
     prefix_match = AllowMatch(prefix_match)
     byte_match = AllowMatch(byte_match)

mergekit/tokensurgeon/token_basis.py

@@ -10,29 +10,29 @@
 
 
 def sparse_linear_basis(
-    embeddings: torch.Tensor,
+    points: torch.Tensor,
     k: int,
     d: int,
     eps: float = 1e-8,
 ) -> Tuple[torch.LongTensor, torch.Tensor]:
     """
-    Form an approximate orthogonal basis from sparse linear combinations of the input embeddings.
+    Form an approximate orthogonal basis from sparse linear combinations of input points.
     Args:
-        embeddings: (num_pts, embed_dim) tensor of embeddings
+        points: (num_pts, embed_dim) tensor of input points
         k: number of points to select per basis vector
         d: dimensionality of the basis
         eps: numerical stability parameter
     Returns:
         indices: (d, k) tensor of selected indices
         coeffs: (d, k) tensor of coefficients for each selected point
     """
-    assert embeddings.dim() == 2
-    num_pts, embed_dim = embeddings.shape
+    assert points.dim() == 2
+    num_pts, embed_dim = points.shape
     assert k <= num_pts, "k must be less than or equal to the number of points"
     assert d <= embed_dim, "d must be less than or equal to the embedding dimension"
 
-    mean_embed = embeddings.mean(dim=0)
-    centered_embeddings = (embeddings - mean_embed).to(torch.float32)
+    mean_embed = points.mean(dim=0)
+    centered_embeddings = (points - mean_embed).to(torch.float32)
     covariance_matrix = (
         centered_embeddings.T @ centered_embeddings
     ) / num_pts  # (embed_dim, embed_dim)