Busca em feixe (Beam Search)

Visão geral

Busca em feixe (beam search) é um algoritmo amplamente usado de decodificação heurística (heuristic decoding) para modelos de sequência (e.g., modelos codificador–decodificador (encoder–decoder models), decodificadores da Arquitetura Transformer, reconhecedores de fala CTC/RNN-T). Ele gera uma sequência de saída token por token, mantendo um conjunto de tamanho fixo das top‑k hipóteses parciais (o “feixe”) de acordo com uma função de pontuação — tipicamente a log-probabilidade cumulativa da sequência sob o modelo (opcionalmente com normalização e penalidades).

A busca em feixe fica entre:

  • Decodificação Gulosa: mantém apenas a melhor hipótese parcial (k=1)
  • Busca exaustiva: mantém todas as hipóteses (inviável para vocabulários reais e comprimentos de sequência)
  • Decodificação baseada em amostragem (sampling-based decoding) como Amostragem Top-k ou Amostragem por Núcleo: troca determinismo por diversidade

Na prática, a busca em feixe é mais valiosa quando você quer gerações de alta verossimilhança e estáveis (e.g., tradução automática, ASR) e pode arcar com computação extra no momento da inferência.

O problema de decodificação (o que a busca em feixe está aproximando)

Muitos modelos de sequência definem uma probabilidade sobre sequências de tokens:

[ P(y_{1:T}\mid x) = \prod_{t=1}^{T} P(y_t \mid y_{<t}, x) ]

onde:

  • (x) é uma entrada de condicionamento opcional (prompt, frase de origem, características de áudio, etc.)
  • (y_t) é o token no passo (t)
  • (T) é o comprimento de saída (muitas vezes desconhecido de antemão; o token EOS encerra a sequência)

Objetivo da decodificação: encontrar a sequência (y) com probabilidade máxima:

[ y^* = \arg\max_y \log P(y\mid x) ]

A busca exata é exponencial em (T) e no tamanho do vocabulário (V). A busca em feixe é uma aproximação pragmática que poda cedo sequências parciais com baixa pontuação.

Ideia central: manter as melhores k hipóteses parciais

Em cada passo de decodificação, a busca em feixe:

  1. Mantém um conjunto de até k sequências parciais (o feixe).
  2. Expande cada sequência parcial em um token (muitas vezes sobre muitos possíveis próximos tokens).
  3. Atribui pontuação aos candidatos resultantes.
  4. Mantém apenas os top k candidatos por pontuação.
  5. Repete até a terminação (e.g., sequências completas suficientes ou comprimento máximo atingido).

Como ela descarta a maioria dos candidatos, não há garantia de encontrar a sequência globalmente ótima a menos que (k) seja extremamente grande (aproximando-se da busca exaustiva).

Pontuação no espaço logarítmico

A busca em feixe quase sempre usa log-probabilidades por estabilidade numérica:

[ \text{score}(y_{1:t}) = \sum_{i=1}^{t} \log P(y_i\mid y_{<i}, x) ]

Usar somas de log-probabilidades também torna “adicionar mais um token” naturalmente incremental.

Algoritmo básico de busca em feixe (geração de texto)

Abaixo está um esqueleto típico de busca em feixe para um decodificador autorregressivo (e.g., um LM Transformer). Implementações reais adicionam processamento em lote, cache e restrições.

def beam_search(model, x, beam_width=4, max_len=50, eos_id=2, length_penalty_alpha=0.0):
    # Each hypothesis: (tokens, logp, finished)
    beam = [([model.bos_id], 0.0, False)]

    for t in range(max_len):
        candidates = []

        for tokens, logp, finished in beam:
            if finished:
                candidates.append((tokens, logp, True))
                continue

            next_logprobs = model.next_token_logprobs(x, tokens)  # shape: [V]
            # Often you only take top-M next tokens for efficiency (M >= beam_width)
            top_tokens = topk_indices(next_logprobs, k=beam_width)

            for tok in top_tokens:
                new_tokens = tokens + [tok]
                new_logp = logp + float(next_logprobs[tok])
                new_finished = (tok == eos_id)
                candidates.append((new_tokens, new_logp, new_finished))

        # Rank candidates by (optionally normalized) score
        def norm_score(h):
            tokens, logp, finished = h
            length = max(1, len(tokens) - 1)  # exclude BOS
            if length_penalty_alpha == 0.0:
                return logp
            # GNMT-style length penalty
            lp = ((5 + length) / 6) ** length_penalty_alpha
            return logp / lp

        candidates.sort(key=norm_score, reverse=True)
        beam = candidates[:beam_width]

        # Optional early stopping if all hypotheses finished
        if all(f for _, _, f in beam):
            break

    # Return best hypothesis (often best finished one)
    best = max(beam, key=lambda h: norm_score(h))
    return best[0]

Complexidade de tempo/espaço (intuição)

Se você expande cada hipótese do feixe por (M) candidatos de próximo token (frequentemente (M = k) ou ligeiramente maior), então cada passo lida com cerca de (k \cdot M) candidatos. A parte cara é computar as probabilidades do modelo, o que é aproximadamente:

  • Sem cache: (O(k \cdot T)) passagens forward (muito caro)
  • Com cache (comum em Transformers): passos incrementais reduzem computação repetida

Na prática, a busca em feixe costuma ser k vezes mais cara do que a decodificação gulosa (mais overhead).

Viés de comprimento e normalização de comprimento (length normalization)

Um problema prático importante: como log-probabilidades são negativas (já que probabilidades estão em ((0,1])), a soma (\sum \log p) tende a ficar mais negativa conforme as sequências ficam mais longas. Isso pode fazer a busca em feixe preferir saídas curtas (encerrando cedo com EOS), especialmente quando EOS é relativamente provável.

Correções comuns:

1) Log-probabilidade média

[ \text{score}\text{avg}(y{1:T}) = \frac{1}{T} \sum_{t=1}^{T} \log P(y_t\mid y_{<t}, x) ]

Isso frequentemente reduz o viés por saídas curtas, mas pode favorecer demais sequências muito longas em alguns cenários.

2) Penalidade de comprimento GNMT (GNMT length penalty) (amplamente usada)

Uma forma popular (do Google NMT) é:

[ \text{LP}(T) = \left(\frac{5+T}{6}\right)^\alpha,\quad \text{score} = \frac{\log P(y)}{\text{LP}(T)} ]

  • (\alpha = 0): sem normalização
  • (\alpha) maior: mais incentivo para sequências longas

3) Restrições de comprimento mínimo

Para tarefas de geração, você pode exigir pelo menos min_len tokens antes que EOS seja permitido, evitando encerramento precoce patológico.

4) Penalidades de cobertura (coverage) e de repetição (repetition) (dependente da tarefa/modelo)

Para tarefas codificador–decodificador como tradução/sumarização, alguns decodificadores adicionam penalidades que incentivam “cobrir” a fonte ou desencorajam frases repetidas. Sistemas Transformer modernos frequentemente dependem mais de treinamento e reclassificação do que de penalidades artesanais, mas elas ainda aparecem em sistemas de produção.

Exemplo prático simples (por que a busca em feixe ajuda)

Suponha que um modelo minúsculo gere uma frase de 2 tokens após BOS:

Vocabulário: {A, B, EOS}

No passo 1:

  • (P(A)=0.6), (P(B)=0.4)

No passo 2:

  • Se o token anterior foi A: (P(EOS)=0.3), (P(A)=0.2), (P(B)=0.5)
  • Se o token anterior foi B: (P(EOS)=0.9), (P(A)=0.05), (P(B)=0.05)

Decodificação gulosa (k=1):

  • Passo 1 escolhe A (0.6)
  • Passo 2 escolhe B (0.5)
  • Resultado: A B com prob (0.6 \cdot 0.5 = 0.30)

Busca em feixe (k=2) mantém ambos os candidatos após o passo 1: [A] e [B].

No passo 2:

  • De [A], a melhor extensão pode ser A B com prob 0.30
  • De [B], a melhor extensão é B EOS com prob (0.4 \cdot 0.9 = 0.36)

A busca em feixe encontra B EOS, o que a decodificação gulosa perde porque se comprometeu cedo demais.

Isso ilustra o valor central da busca em feixe: evitar escolhas locais irreversíveis mantendo múltiplas hipóteses parciais promissoras.

Busca em feixe na geração de texto com LLMs (large language models)

A busca em feixe é comum em tradução automática e em alguns cenários de sumarização porque:

  • As saídas devem ser de alta verossimilhança e consistentes
  • Determinismo é frequentemente desejável
  • Em geral há uma região relativamente estreita de saídas “corretas”

No entanto, para geração aberta (open-ended) (escrita criativa, diálogo), a busca em feixe pode ter desempenho ruim:

  • Ela frequentemente produz respostas genéricas ou “seguras”
  • Aumentar a largura do feixe pode reduzir a diversidade e às vezes a qualidade (a “maldição da busca em feixe (beam search curse)” observada em diálogo e em algumas configurações de sumarização)
  • Ela pode amplificar padrões repetitivos se a massa de probabilidade do modelo tiver esse formato

Como resultado, muitas aplicações de LLM preferem decodificação baseada em amostragem (veja Amostragem Top-k, Amostragem por Núcleo e Temperatura) por diversidade e variação mais semelhante à humana, frequentemente com filtros de segurança e reclassificadores.

Hiperparâmetros comuns de geração de texto

  • Largura do feixe (beam width) (k): k maior → mais busca, decodificação mais lenta; frequentemente 2–8 para muitas tarefas
  • Comprimento máximo / comprimento mínimo
  • Penalidade de comprimento (length penalty) (\alpha): crítica para tradução/sumarização
  • Parada antecipada (early stopping): parar quando a melhor hipótese finalizada é melhor do que todas as não finalizadas (cuidado: depende da normalização)
  • No-repeat n-gram / penalidades de repetição: correções pragmáticas para problemas de repetição
  • Retornar lista n-best (n-best list): manter as top N sequências completas para reclassificação ou estimativa de incerteza

Busca em feixe para CTC (Classificação Temporal Conexionista) (Connectionist Temporal Classification)

CTC é usada em reconhecimento de fala e outras rotulações de sequência sem alinhamento. O modelo produz uma distribuição sobre tokens em cada quadro de tempo, incluindo um símbolo especial em branco (blank). Muitos caminhos diferentes em nível de quadro podem corresponder ao mesmo texto final após colapsar repetições e remover em branco.

Por causa desse mapeamento muitos-para-um, a decodificação não é a mesma coisa que uma geração simples da esquerda para a direita. A abordagem padrão é a busca em feixe por prefixos do CTC (CTC prefix beam search), que mantém probabilidades de prefixos ao considerar se o prefixo termina em branco ou não-em-branco.

Ideia-chave: probabilidades de prefixo e mesclagem

Em CTC, múltiplos alinhamentos produzem o mesmo prefixo de texto. A busca em feixe por prefixos:

  • Acompanha cada prefixo de texto (p) com duas pontuações:
    • (p_b): probabilidade do prefixo terminar com branco
    • (p_{nb}): probabilidade do prefixo terminar com um token não-em-branco
  • Ao estender prefixos, ela mescla (merge) caminhos que resultam no mesmo prefixo, o que é essencial para correção e eficiência.

Essa mesclagem é o motivo de a busca em feixe para CTC ser mais especializada do que o algoritmo genérico mostrado antes.

Adicionando um modelo de linguagem externo (fusão superficial) (shallow fusion)

Sistemas de ASR frequentemente combinam pontuações do modelo acústico e do modelo de linguagem durante a busca em feixe:

[ \text{score} = \log P_\text{CTC}(y \mid x) + \lambda \log P_\text{LM}(y) + \beta \cdot |y| ]

  • (\lambda): peso do modelo de linguagem
  • (\beta): bônus de inserção / termo de viés de comprimento

Esta é uma técnica prática e poderosa porque o modelo acústico pode estar incerto, e o modelo de linguagem ajuda a escolher saídas linguisticamente plausíveis.

Busca em feixe para RNN-T (Transdutor de Rede Neural Recorrente) (Recurrent Neural Network Transducer)

RNN-T é comum em ASR em streaming. Diferentemente de CTC, RNN-T modela explicitamente interações entre:

  • Um codificador (encoder) sobre quadros acústicos
  • Uma rede de predição (prediction network) sobre tokens de saída anteriores
  • Uma rede conjunta (joint network) que produz distribuições sobre o próximo token + branco

A decodificação envolve decidir quando emitir um token versus emitir branco para avançar no tempo. A busca em feixe para RNN-T tipicamente:

  • Mantém um feixe de hipóteses sobre históricos de tokens de saída
  • Em cada quadro acústico, executa um laço interno limitado de expansões de tokens até que branco seja emitido (ou um limite seja atingido)
  • Usa limiares de poda e limites máximos de expansão para manter o tempo de execução limitado

A busca em feixe para RNN-T é, portanto, uma busca bidimensional sobre tempo e tokens de saída, exigindo heurísticas agressivas para desempenho em tempo real.

Variantes comuns de busca em feixe

A busca em feixe é uma estrutura; muitas variantes existem para lidar com seus problemas conhecidos.

Busca em feixe diversa (Diverse Beam Search, DBS)

A busca em feixe padrão tende a produzir hipóteses muito semelhantes dentro do feixe (pequenas variações do mesmo prefixo). A busca em feixe diversa incentiva diversidade ao particionar o feixe em grupos e adicionar uma penalidade de diversidade para que cada grupo explore regiões diferentes.

Uma abordagem comum:

  • Dividir o feixe em (G) grupos de tamanho (k/G)
  • Decodificar grupos sequencialmente em cada passo
  • Penalizar tokens que já foram escolhidos por grupos anteriores (e.g., diversidade de Hamming)

Isso pode melhorar resultados quando você quer múltiplos candidatos distintos (e.g., para reclassificação, ou para produzir traduções alternativas).

Busca em feixe estocástica / amostragem em feixe (beam sampling)

Em vez de sempre manter os top-k candidatos determinísticos, você amostra candidatos proporcionalmente às suas pontuações (possivelmente com temperatura). Isso combina exploração (amostragem) com rastreamento de múltiplas hipóteses (feixe), às vezes melhorando a diversidade enquanto preserva alguma estrutura de busca.

Busca em feixe com restrições (constrained beam search) (restrições lexicais)

Às vezes as saídas precisam incluir frases específicas (e.g., entidades nomeadas, terminologia). A decodificação com restrições modifica a busca em feixe para que as hipóteses acompanhem estados de satisfação de restrições (frequentemente via máquinas de estados finitos), garantindo que a saída final cumpra as restrições.

Isso é comum em geração controlada, data-to-text e tradução sensível à terminologia.

Reclassificação (reranking) e risco mínimo de Bayes (minimum Bayes risk, MBR)

A busca em feixe frequentemente retorna uma lista n-best. Um modelo ou critério de segunda etapa pode reclassificar:

  • Um modelo mais forte (ou um comitê (ensemble))
  • Um modelo de recompensa (comum em pipelines de pós-treinamento de LLM; veja Aprendizado por Reforço a partir de Feedback Humano)
  • Decodificação MBR: escolher o candidato que minimiza a perda esperada sob a distribuição do modelo (popular em pipelines de tradução e ASR)

Em sistemas modernos, “busca em feixe + reclassificação” é frequentemente mais eficaz do que larguras de feixe extremamente grandes.

Escolhendo largura do feixe e outros hiperparâmetros (orientação prática)

Largura do feixe (k)

  • k=1: decodificação gulosa, mais rápida
  • k=2–8: ponto de equilíbrio comum em muitas tarefas seq2seq
  • k>10: frequentemente retornos decrescentes; pode degradar a qualidade em algumas tarefas e aumenta a latência

Normalização / penalidade de comprimento

  • Essencial para tradução/sumarização
  • Ajuste (\alpha) em um conjunto de validação com sua métrica (BLEU, chrF, ROUGE, impacto no WER indiretamente via inserções/deleções em ASR)

Parada antecipada

Duas abordagens comuns:

  • Parar quando a melhor hipótese no feixe estiver finalizada (rápido, mas pode estar errado sob algumas normalizações)
  • Parar quando você tiver coletado num_return_sequences hipóteses finalizadas e todas as não finalizadas não puderem superar a melhor pontuação finalizada (mais fundamentado, mais computação)

Para ASR (CTC/RNN-T)

  • Ajuste o peso do modelo de linguagem (\lambda) e o bônus de inserção (\beta)
  • Use limiares de poda para controlar o fator de tempo real
  • Considere incompatibilidade de domínio: um modelo de linguagem externo forte pode ajudar muito, mas também pode introduzir viés se estiver desalinhado

Limitações e modos de falha

  • Erros de busca devido à poda: a melhor sequência global pode cair fora do feixe cedo.
  • Baixa diversidade: feixes colapsam para candidatos quase idênticos, especialmente com distribuições muito concentradas.
  • Patologias de comprimento: sem normalização, a busca em feixe pode preferir saídas curtas; com penalidade de comprimento forte demais, pode preferir saídas longas demais.
  • Desalinhamento com preferências humanas: maximizar verossimilhança nem sempre maximiza qualidade percebida (notavelmente em diálogo). Por isso, amostragem ou reclassificação é comum em aplicações de LLM.

Quando usar busca em feixe vs alternativas

Use busca em feixe quando:

  • Você quer decodificação determinística, de alta verossimilhança
  • Você tem uma noção relativamente bem definida de correção (tradução, ASR)
  • Você pode pagar o custo extra de inferência

Prefira métodos de amostragem quando:

Resumo

A busca em feixe é uma estratégia prática e de uso geral de decodificação que aproxima a geração de sequências de probabilidade máxima ao manter as hipóteses parciais top‑k em cada passo. Sua eficácia depende fortemente da largura do feixe, da normalização de comprimento e de pontuações específicas da tarefa (especialmente em ASR com CTC e RNN-T, onde tratamento especializado de prefixos/branco e fusão opcional com modelo de linguagem são padrão). Variantes como busca em feixe diversa e pipelines de reclassificação lidam com problemas comuns como baixa diversidade e desalinhamento com preferências, tornando a busca em feixe uma ferramenta central na inferência moderna de modelos de sequência.