Use max_active_dims for pooled latent pruning

README.md

@@ -207,8 +207,8 @@ documents = [
     "Wind turbines convert the kinetic energy of wind into electricity.",
 ]
 
-query_embeddings = model.encode_query(queries, top_k=64)
-document_embeddings = model.encode_document(documents, top_k=512)
+query_embeddings = model.encode_query(queries, max_active_dims=64)
+document_embeddings = model.encode_document(documents, max_active_dims=512)
 
 print(query_embeddings.shape, document_embeddings.shape)
 # torch.Size([1, 32768]) torch.Size([2, 32768])
@@ -221,16 +221,18 @@ print(scores)
 
 ### Choosing Query and Document Budgets
 
-Use `top_k` independently on each encoding call to cap the number of nonzero features **after pooling**. This is separate from the per-token top-32 operation and does not change the vector dimensionality.
+Set `max_active_dims=N` to keep the **top-N latents with the highest weights** in each pooled vector and set the rest to zero. You can choose a different N for queries and documents. This is separate from the per-token top-32 operation and does not change the vector dimensionality.
+
+Use a positive integer for a cap, or `None` for the saved unpruned defaults; `0` is invalid.
 
 To reproduce the unpruned setting:
 
 ```python
-query_embeddings = model.encode_query(queries, top_k=0)
-document_embeddings = model.encode_document(documents, top_k=0)
+query_embeddings = model.encode_query(queries, max_active_dims=None)
+document_embeddings = model.encode_document(documents, max_active_dims=None)
 ```
 
-Omitting `top_k` also uses the saved unpruned defaults. To use the **q-full / d-512** setting, keep the query call at `top_k=0` and set the document call to `top_k=512`.
+Omitting `max_active_dims` also uses the saved unpruned defaults. To use the **q-full / d-512** setting, keep the query call at `max_active_dims=None` and set the document call to `max_active_dims=512`.
 
 Dimensions represent **learned latents**, so they do not map directly to tokenizer vocabulary entries. Inspect the sparse indices and weights; `SparseEncoder.decode()` cannot label them as vocabulary tokens. Retrieval uses raw dot products, so leave final vector normalization disabled.
 

config.json

@@ -94,7 +94,7 @@
   },
   "inference": {
     "document_length": 8192,
-    "document_top_k": 0,
+    "document_max_active_dims": null,
     "head_chunk_size": 2048,
     "pooling": "max_relu",
     "prompts": {
@@ -102,7 +102,7 @@
       "query": "query: "
     },
     "query_length": 8192,
-    "query_top_k": 0,
+    "query_max_active_dims": null,
     "score": "dot",
     "token_k": 32
   },

modeling_slideon.py

@@ -22,7 +22,7 @@ class SlideOnConfig(PretrainedConfig):
         self.projection_configs = projection_configs or []
         self.head_config = head_config or dict(d_model=768, n_latents=32768, d_lift=0)
         self.inference = inference or dict(
-            token_k=32, query_top_k=0, document_top_k=0,
+            token_k=32, query_max_active_dims=None, document_max_active_dims=None,
             query_length=8192, document_length=8192, head_chunk_size=2048,
             prompts={"query": "", "document": ""}, pooling="max_relu", score="dot")
         self.tokenizer_settings = tokenizer_settings or {}
@@ -134,7 +134,9 @@ class SlideOnModel(PreTrainedModel):
 
 
     def forward(self, input_ids, attention_mask, token_type_ids=None, task=None,
-                slideon_task=None, top_k=None, token_k=None, **kwargs):
+                slideon_task=None, max_active_dims=None, token_k=None, **kwargs):
+        if "top_k" in kwargs:
+            raise TypeError("top_k has been removed; use max_active_dims (None for unpruned)")
         features = {"input_ids": input_ids, "attention_mask": attention_mask}
         if slideon_task is not None or task is not None:
             features["slideon_task"] = slideon_task or task
@@ -144,10 +146,10 @@ class SlideOnModel(PreTrainedModel):
         if features.get("slideon_task", task) != task:
             raise ValueError("Preprocessing and forward tasks must match")
         n_latents = self.get_embedding_dimension()
-        cap = self.settings[task + "_top_k"] if top_k is None else top_k
+        cap = self.settings.get(task + "_max_active_dims") if max_active_dims is None else max_active_dims
         k = self.settings["token_k"] if token_k is None else token_k
-        if not isinstance(cap, int) or not 0 <= cap <= n_latents:
-            raise ValueError(f"top_k must be an integer between 0 and {n_latents}")
+        if cap is not None and (type(cap) is not int or cap <= 0):
+            raise ValueError("max_active_dims must be a positive integer or None")
         if not isinstance(k, int) or not 1 <= k <= n_latents:
             raise ValueError(f"token_k must be an integer between 1 and {n_latents}")
         model_inputs = {
@@ -199,8 +201,9 @@ class SlideOnModel(PreTrainedModel):
             pooled.view(mask.shape[0], n_latents)
             * self.head.gate_relu_phi.float().relu()
         )
-        if cap:
-            values, indices = pooled.topk(cap, dim=-1)
+        if cap is not None:
+            # Prune before FP16 rounding
+            values, indices = pooled.topk(min(cap, n_latents), dim=-1)
             pooled = torch.zeros_like(pooled).scatter(1, indices, values)
         # The reference stores retrieval weights in FP16; retain that rounding.
         features["sentence_embedding"] = pooled if self.training else pooled.half().float()

slideon.py

@@ -8,7 +8,7 @@ from transformers import AutoConfig, AutoModel, AutoTokenizer
 
 class SlideOn(InputModule):
     config_file_name = "config.json"
-    forward_kwargs = {"task", "top_k", "token_k"}
+    forward_kwargs = {"task", "max_active_dims", "token_k"}
     save_in_root = True
 
     def __init__(self, checkpoint, model_kwargs=None, processor_kwargs=None,
@@ -75,10 +75,12 @@ class SlideOn(InputModule):
         features["slideon_task"] = task
         return features
 
-    def forward(self, features, task=None, top_k=None, token_k=None, **kwargs):
+    def forward(self, features, task=None, max_active_dims=None, token_k=None, **kwargs):
+        if "top_k" in kwargs:
+            raise TypeError("top_k has been removed; use max_active_dims (None for unpruned)")
         inputs = {key: features[key] for key in
                   ("input_ids", "attention_mask", "token_type_ids", "slideon_task") if key in features}
-        features.update(self.model(**inputs, task=task, top_k=top_k, token_k=token_k))
+        features.update(self.model(**inputs, task=task, max_active_dims=max_active_dims, token_k=token_k))
         return features
 
     def save(self, output_path, safe_serialization=True, **kwargs):