Skip to content

ensemble

Ensemble

Bases: Potential

An ensemble potential combining two language models using a weighted operation.

The ensemble's weight of a sequence combines the two potentials' weights of it with a weighted operation (e.g., weighted arithmetic or geometric mean, min, max): prefix(x) = op(p1.prefix(x), p2.prefix(x)), and likewise for complete. logw_next is consistent with prefix, so sampling tokens from it (e.g. with direct_token_sampler) targets this combined distribution over sequences.

Parameters:

Name Type Description Default
p1 Potential

First potential (language model)

 required
p2 Potential

Second potential (language model)

 required
op str | float

"sum", "prod", "harmonic", "min", "max", or a power-mean exponent p (see convert_to_weighted_logop)

 required
a float

Weighting parameter between 0 and 1 (default 0.5 for equal weighting). When a=0.5, models are weighted equally. For a != 0.5, the combination is weighted: a * model1 + (1-a) * model2

 0.5
cache_size int

Number of contexts whose next-token weights are kept, so that the potentials' prefix weights of their extensions need no recomputation. Set it to at least the number of particles. Defaults to 256.

 256

Attributes:

Name Type Description
p1

First potential
p2

Second potential
op

Weighted log operation function
p1_vocab_idxs

Indices mapping unified vocabulary to p1's vocabulary
p2_vocab_idxs

Indices mapping unified vocabulary to p2's vocabulary
Example 
from genlm.control import PromptedLLM, Ensemble

# Create two language model potentials
p1 = PromptedLLM.from_name("openai-community/gpt2")
p2 = PromptedLLM.from_name("openai-community/gpt2")

# Create an ensemble with weighted geometric mean (a=0.5)
ensemble = Ensemble(p1, p2, op="prod", a=0.5)

# Sample from the ensemble with SMC
sequences = await direct_token_sampler(ensemble).smc(
    n_particles=10, ess_threshold=0.5, max_tokens=20
)
Note

Both potentials must have the same vocabulary (typically the same tokenizer). To ensemble models with different tokenizers, ensemble them at the byte level with ByteLLM potentials.

Source code in genlm/control/potential/built_in/ensemble.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
class Ensemble(Potential):
    """An ensemble potential combining two language models using a weighted operation.

    The ensemble's weight of a sequence combines the two potentials' weights of it
    with a weighted operation (e.g., weighted arithmetic or geometric mean, min, max):
    `prefix(x) = op(p1.prefix(x), p2.prefix(x))`, and likewise for `complete`.
    `logw_next` is consistent with `prefix`, so sampling tokens from it (e.g. with
    `direct_token_sampler`) targets this combined distribution over sequences.

    Args:
        p1 (Potential): First potential (language model)
        p2 (Potential): Second potential (language model)
        op (str | float): "sum", "prod", "harmonic", "min", "max", or a power-mean
            exponent p (see `convert_to_weighted_logop`)
        a (float): Weighting parameter between 0 and 1 (default 0.5 for equal weighting).
            When a=0.5, models are weighted equally. For a != 0.5, the combination
            is weighted: a * model1 + (1-a) * model2
        cache_size (int): Number of contexts whose next-token weights are kept, so that
            the potentials' prefix weights of their extensions need no recomputation.
            Set it to at least the number of particles. Defaults to 256.

    Attributes:
        p1: First potential
        p2: Second potential
        op: Weighted log operation function
        p1_vocab_idxs: Indices mapping unified vocabulary to p1's vocabulary
        p2_vocab_idxs: Indices mapping unified vocabulary to p2's vocabulary

    Example:
        ```python
        from genlm.control import PromptedLLM, Ensemble

        # Create two language model potentials
        p1 = PromptedLLM.from_name("openai-community/gpt2")
        p2 = PromptedLLM.from_name("openai-community/gpt2")

        # Create an ensemble with weighted geometric mean (a=0.5)
        ensemble = Ensemble(p1, p2, op="prod", a=0.5)

        # Sample from the ensemble with SMC
        sequences = await direct_token_sampler(ensemble).smc(
            n_particles=10, ess_threshold=0.5, max_tokens=20
        )
        ```

    Note:
        Both potentials must have the same vocabulary (typically the same tokenizer).
        To ensemble models with different tokenizers, ensemble them at the byte level
        with `ByteLLM` potentials.
    """

    def __init__(
        self,
        p1: Potential,
        p2: Potential,
        op: Union[str, float],
        a: float = 0.5,
        cache_size: int = 256,
    ):
        self.p1 = p1
        self.p2 = p2
        self.op = convert_to_weighted_logop(op, a)

        if set(p1.vocab) != set(p2.vocab):
            raise ValueError(
                "Ensemble requires potentials with the same vocabulary. "
                "To ensemble models with different tokenizers, ensemble ByteLLM potentials."
            )

        vocab = list(dict.fromkeys(p1.vocab + p2.vocab))
        super().__init__(vocabulary=vocab)

        self.p1_vocab_idxs = [self.p1.lookup[x] for x in self.vocab_eos]
        self.p2_vocab_idxs = [self.p2.lookup[x] for x in self.vocab_eos]

        # context -> (p1 prefix, p2 prefix, p1 next-token row, p2 next-token row)
        self._rows = LRUCache(cache_size)

    async def _component_prefixes(self, contexts):
        """Each potential's prefix log weight of each context.

        Read off the cached next-token weights of the context's parent when available,
        and computed with the potentials' `prefix` otherwise.

        Returns:
            (np.ndarray): Shape `[N, 2]`, the two potentials' prefix log weights.
        """
        W = np.empty((len(contexts), 2))
        missing = []
        for n, context in enumerate(contexts):
            parent = self._rows.get(tuple(context[:-1])) if context else None
            if parent is not None:
                w1, w2, row1, row2 = parent
                i = self.lookup[context[-1]]
                W[n] = w1 + row1[i], w2 + row2[i]
            else:
                missing.append(n)
        if missing:
            ctxs = [contexts[n] for n in missing]
            W1, W2 = await asyncio.gather(
                self.p1.batch_prefix(ctxs), self.p2.batch_prefix(ctxs)
            )
            W[missing, 0] = to_numpy(W1)
            W[missing, 1] = to_numpy(W2)
        return W

    async def component_logws(self, context: List[str]) -> Tuple[float, float]:
        """Each potential's log weight of `context`.

        For a context ending in EOS these are the potentials' `complete` weights of
        the sequence before it, otherwise their `prefix` weights. Together with the
        ensemble's own weight, e.g. to record per-model log weights of SMC samples.

        Args:
            context (List[str]): The context tokens, optionally ending in EOS.

        Returns:
            Tuple[float, float]: The log weights under `p1` and `p2`.
        """
        if context and context[-1] == self.eos:
            parent = self._rows.get(tuple(context[:-1]))
            if parent is not None:
                w1, w2, row1, row2 = parent
                i = self.lookup[self.eos]
                return float(w1 + row1[i]), float(w2 + row2[i])
            w1, w2 = await asyncio.gather(
                self.p1.complete(context[:-1]), self.p2.complete(context[:-1])
            )
            return float(w1), float(w2)
        ((w1, w2),) = await self._component_prefixes([context])
        return float(w1), float(w2)

    async def prefix(self, context: List[str]) -> float:
        """Compute log weights for the prefix using both potentials.

        Args:
            context (List[str]): The context tokens

        Returns:
            float: Combined log weight from both potentials using the ensemble operation
        """
        ((w1, w2),) = await self._component_prefixes([context])
        return self.op(w1, w2)

    async def complete(self, context: List[str]) -> float:
        """Compute completion log weights using both potentials.

        Args:
            context (List[str]): The context tokens

        Returns:
            float: Combined completion log weight from both potentials
        """
        p1_logw, p2_logw = await asyncio.gather(
            self.p1.complete(context), self.p2.complete(context)
        )
        return self.op(p1_logw, p2_logw)

    async def batch_prefix(self, contexts):
        if not contexts:
            raise ValueError("Contexts must be non-empty.")
        W = await self._component_prefixes(contexts)
        return self.op(W[:, 0], W[:, 1])

    async def batch_complete(self, contexts):
        if not contexts:
            raise ValueError("Contexts must be non-empty.")
        W1, W2 = await asyncio.gather(
            self.p1.batch_complete(contexts), self.p2.batch_complete(contexts)
        )
        return self.op(to_numpy(W1), to_numpy(W2))

    async def logw_next(self, context: List[str]):
        """Next-token log weights, `prefix(context + [x]) - prefix(context)`.

        Args:
            context (List[str]): The context tokens

        Returns:
            (LazyWeights): Log weights over `self.vocab_eos`.
        """
        batch = await self.batch_logw_next([context])
        return batch.spawn(batch.weights[0])

    async def batch_logw_next(self, contexts: List[List[str]]):
        """Batched version of logw_next for Ensemble.

        Args:
            contexts (List[List[str]]): List of context token sequences

        Returns:
            (LazyWeights): Batched log weights, `.weights` of shape `[N, V+1]`, row `n`
                being `prefix(contexts[n] + [x]) - prefix(contexts[n])`.
        """
        (Ws1, Ws2), W = await asyncio.gather(
            asyncio.gather(
                self.p1.batch_logw_next(contexts), self.p2.batch_logw_next(contexts)
            ),
            self._component_prefixes(contexts),
        )
        W1, W2 = W[:, :1], W[:, 1:]
        base = self.op(W1, W2)
        zero = np.isneginf(base[:, 0])
        if zero.any():
            context = contexts[int(np.argmax(zero))]
            raise ValueError(f"Context {context!r} has weight zero under `prefix`.")
        rows1 = to_numpy(Ws1.weights)[:, self.p1_vocab_idxs]
        rows2 = to_numpy(Ws2.weights)[:, self.p2_vocab_idxs]
        for n, context in enumerate(contexts):
            self._rows[tuple(context)] = (W[n, 0], W[n, 1], rows1[n], rows2[n])
        return self.make_lazy_weights(self.op(W1 + rows1, W2 + rows2) - base)

component_logws(context) async

Each potential's log weight of context.

For a context ending in EOS these are the potentials' complete weights of the sequence before it, otherwise their prefix weights. Together with the ensemble's own weight, e.g. to record per-model log weights of SMC samples.

Parameters:

Name Type Description Default
context List[str]

The context tokens, optionally ending in EOS.

 required

Returns:

Type Description
Tuple[float, float]

Tuple[float, float]: The log weights under p1 and p2.
Source code in genlm/control/potential/built_in/ensemble.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
async def component_logws(self, context: List[str]) -> Tuple[float, float]:
    """Each potential's log weight of `context`.

    For a context ending in EOS these are the potentials' `complete` weights of
    the sequence before it, otherwise their `prefix` weights. Together with the
    ensemble's own weight, e.g. to record per-model log weights of SMC samples.

    Args:
        context (List[str]): The context tokens, optionally ending in EOS.

    Returns:
        Tuple[float, float]: The log weights under `p1` and `p2`.
    """
    if context and context[-1] == self.eos:
        parent = self._rows.get(tuple(context[:-1]))
        if parent is not None:
            w1, w2, row1, row2 = parent
            i = self.lookup[self.eos]
            return float(w1 + row1[i]), float(w2 + row2[i])
        w1, w2 = await asyncio.gather(
            self.p1.complete(context[:-1]), self.p2.complete(context[:-1])
        )
        return float(w1), float(w2)
    ((w1, w2),) = await self._component_prefixes([context])
    return float(w1), float(w2)

prefix(context) async

Compute log weights for the prefix using both potentials.

Parameters:

Name Type Description Default
context List[str]

The context tokens

 required

Returns:

Name Type Description
float float

Combined log weight from both potentials using the ensemble operation
Source code in genlm/control/potential/built_in/ensemble.py 
142
143
144
145
146
147
148
149
150
151
152
async def prefix(self, context: List[str]) -> float:
    """Compute log weights for the prefix using both potentials.

    Args:
        context (List[str]): The context tokens

    Returns:
        float: Combined log weight from both potentials using the ensemble operation
    """
    ((w1, w2),) = await self._component_prefixes([context])
    return self.op(w1, w2)

complete(context) async

Compute completion log weights using both potentials.

Parameters:

Name Type Description Default
context List[str]

The context tokens

 required

Returns:

Name Type Description
float float

Combined completion log weight from both potentials
Source code in genlm/control/potential/built_in/ensemble.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
async def complete(self, context: List[str]) -> float:
    """Compute completion log weights using both potentials.

    Args:
        context (List[str]): The context tokens

    Returns:
        float: Combined completion log weight from both potentials
    """
    p1_logw, p2_logw = await asyncio.gather(
        self.p1.complete(context), self.p2.complete(context)
    )
    return self.op(p1_logw, p2_logw)

logw_next(context) async

Next-token log weights, prefix(context + [x]) - prefix(context).

Parameters:

Name Type Description Default
context List[str]

The context tokens

 required

Returns:

Type Description
LazyWeights

Log weights over self.vocab_eos.
Source code in genlm/control/potential/built_in/ensemble.py 
182
183
184
185
186
187
188
189
190
191
192
async def logw_next(self, context: List[str]):
    """Next-token log weights, `prefix(context + [x]) - prefix(context)`.

    Args:
        context (List[str]): The context tokens

    Returns:
        (LazyWeights): Log weights over `self.vocab_eos`.
    """
    batch = await self.batch_logw_next([context])
    return batch.spawn(batch.weights[0])

batch_logw_next(contexts) async

Batched version of logw_next for Ensemble.

Parameters:

Name Type Description Default
contexts List[List[str]]

List of context token sequences

 required

Returns:

Type Description
LazyWeights

Batched log weights, .weights of shape [N, V+1], row n being prefix(contexts[n] + [x]) - prefix(contexts[n]).
Source code in genlm/control/potential/built_in/ensemble.py 
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
async def batch_logw_next(self, contexts: List[List[str]]):
    """Batched version of logw_next for Ensemble.

    Args:
        contexts (List[List[str]]): List of context token sequences

    Returns:
        (LazyWeights): Batched log weights, `.weights` of shape `[N, V+1]`, row `n`
            being `prefix(contexts[n] + [x]) - prefix(contexts[n])`.
    """
    (Ws1, Ws2), W = await asyncio.gather(
        asyncio.gather(
            self.p1.batch_logw_next(contexts), self.p2.batch_logw_next(contexts)
        ),
        self._component_prefixes(contexts),
    )
    W1, W2 = W[:, :1], W[:, 1:]
    base = self.op(W1, W2)
    zero = np.isneginf(base[:, 0])
    if zero.any():
        context = contexts[int(np.argmax(zero))]
        raise ValueError(f"Context {context!r} has weight zero under `prefix`.")
    rows1 = to_numpy(Ws1.weights)[:, self.p1_vocab_idxs]
    rows2 = to_numpy(Ws2.weights)[:, self.p2_vocab_idxs]
    for n, context in enumerate(contexts):
        self._rows[tuple(context)] = (W[n, 0], W[n, 1], rows1[n], rows2[n])
    return self.make_lazy_weights(self.op(W1 + rows1, W2 + rows2) - base)

convert_to_weighted_logop(op, a=0.5)

Convert an operation to its weighted log-space equivalent.

This function takes an operation and a weighting parameter and returns a function that combines two log-probability arrays using the specified weighted operation.

Parameters:

Name Type Description Default
op str | float

The operation: - a number p: the weighted power mean M_p(x, y) = (a * x^p + (1-a) * y^p)^(1/p), with p = 0 the geometric mean - "sum" (arithmetic mean), "prod" (geometric mean), "harmonic": the power means with p = 1, 0 and -1 - "min", "max": weighted extrema

 required
a float

Weighting parameter between 0 and 1. When a=0.5, equal weighting. For weighted operations: a * model1 + (1-a) * model2

 0.5

Returns:

Type Description
Callable[[ndarray, ndarray], ndarray]

Callable[[np.ndarray, np.ndarray], np.ndarray]: A function that takes two log-probability arrays and returns their weighted combination in log space.

Raises:

Type Description
ValueError

If a is not between 0 and 1, or if op is not recognized.

Examples:

>>> op_func = convert_to_weighted_logop("sum", a=0.5)
>>> x = np.log(np.array([0.3, 0.7]))
>>> y = np.log(np.array([0.6, 0.4]))
>>> result = op_func(x, y)  # Weighted arithmetic mean in log space
>>> result = convert_to_weighted_logop(2.5, a=0.5)(x, y)  # Power mean, p=2.5
Source code in genlm/control/potential/built_in/ensemble.py 
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
def convert_to_weighted_logop(
    op: Union[str, float],
    a: float = 0.5,
) -> Callable[[np.ndarray, np.ndarray], np.ndarray]:
    """Convert an operation to its weighted log-space equivalent.

    This function takes an operation and a weighting parameter and returns
    a function that combines two log-probability arrays using the specified
    weighted operation.

    Args:
        op (str | float): The operation:
            - a number p: the weighted power mean
              M_p(x, y) = (a * x^p + (1-a) * y^p)^(1/p), with p = 0 the geometric mean
            - "sum" (arithmetic mean), "prod" (geometric mean), "harmonic": the
              power means with p = 1, 0 and -1
            - "min", "max": weighted extrema
        a (float): Weighting parameter between 0 and 1. When a=0.5, equal weighting.
            For weighted operations: a * model1 + (1-a) * model2

    Returns:
        Callable[[np.ndarray, np.ndarray], np.ndarray]: A function that takes two
            log-probability arrays and returns their weighted combination in log space.

    Raises:
        ValueError: If a is not between 0 and 1, or if op is not recognized.

    Examples:
        >>> op_func = convert_to_weighted_logop("sum", a=0.5)
        >>> x = np.log(np.array([0.3, 0.7]))
        >>> y = np.log(np.array([0.6, 0.4]))
        >>> result = op_func(x, y)  # Weighted arithmetic mean in log space
        >>> result = convert_to_weighted_logop(2.5, a=0.5)(x, y)  # Power mean, p=2.5
    """
    if not 0 < a < 1:
        raise ValueError("variable a should be between 0 and 1")

    if op == "min":
        return _neginf_for_nan(_weighted_extremum(np.minimum, a))
    if op == "max":
        return _neginf_for_nan(_weighted_extremum(np.maximum, a))

    p = _NAMED_POWERS.get(op, op) if isinstance(op, str) else op
    if not isinstance(p, numbers.Real) or isinstance(p, bool):
        valid = ", ".join(repr(o) for o in [*_NAMED_POWERS, "min", "max"])
        raise ValueError(
            f"Invalid operation: {op!r}. Must be one of {valid}, or a number p "
            "for the power mean."
        )
    return _neginf_for_nan(_power_mean(float(p), a))